diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 23040587a5..d0cd8bf8ef 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -16,6 +16,7 @@ import asyncio import atexit +import base64 import collections.abc from concurrent.futures import Future as ConcurrentFuture from concurrent.futures import ThreadPoolExecutor @@ -23,14 +24,19 @@ import dataclasses from dataclasses import dataclass from dataclasses import field +from datetime import date from datetime import datetime +from datetime import timedelta from datetime import timezone +import decimal +import enum import functools import json import logging import math import mimetypes import os +import pathlib import traceback as traceback_module # Enable gRPC fork support so child processes created via os.fork() @@ -53,6 +59,11 @@ from typing import ParamSpec from typing import TYPE_CHECKING from typing import TypeVar +from urllib.parse import parse_qsl +from urllib.parse import quote +from urllib.parse import urlencode +from urllib.parse import urlsplit +from urllib.parse import urlunsplit import uuid import weakref @@ -414,11 +425,218 @@ def _extract_tool_declarations( "id_token", "api_key", "password", + "private_key", + "proxy_authorization", + "google_access_id", + "sig", + "signature", + "token", + "secret", + "authorization", + "x_api_key", + "x_amz_credential", + "x_amz_signature", + "x_goog_credential", + "x_goog_security_token", + "x_goog_signature", }) +# Credentials commonly carried in signed URLs and HTTP error text. These are +# values rather than structured mapping keys in those two surfaces, so they +# need an explicit bounded text/URI pass in addition to _SENSITIVE_KEYS. +_SENSITIVE_TEXT_KEYS = _SENSITIVE_KEYS +_SENSITIVE_TEXT_KEY_RE = re.compile( + r"(?i)(?P(?" + r'"(?:\\.|[^"\\])*"' + r"|'(?:\\.|[^'\\])*'" + r"|[^\s,;&}]+)" +) +_AUTH_HEADER_RE = re.compile( + r"(?im)(?P\b(?:authorization|proxy-authorization|x-api-key|api-key)" + r"[ \t]*:[ \t]*)[^\r\n]*" +) +_BEARER_TOKEN_RE = re.compile( + r"(?i)(?P\bbearer\s+)(?!of(?:\s|$))[^\s,;]+" +) +_BASIC_TOKEN_RE = re.compile( + r"(?i)(?P\bbasic\s+)(?P[A-Za-z0-9+/]+={0,2})" + r"(?![A-Za-z0-9+/=])" +) +_MAX_BASIC_AUTH_TOKEN_CHARS = 16 * 1024 +_ASCII_HEX_DIGITS = frozenset("0123456789abcdefABCDEF") + + +def _is_sensitive_text_key(text: str) -> bool: + """Returns whether ``text`` is exactly a credential-bearing key.""" + normalized = text.lower().replace("-", "_") + return normalized in _SENSITIVE_TEXT_KEYS or normalized.startswith("temp:") + + +def _canonicalize_common_ascii_escapes(text: str) -> str: + """Decodes nested ASCII ``\\u``, ``\\x``, and percent escapes in O(n). + + This is a detection-only representation: callers retain the original text + unless the canonical form exposes a credential construct. The output stack + also handles encodings that reveal another escape introducer, such as + ``%255F`` and ``\\u005cu005f``, without rescanning the whole input. + """ + output: list[str] = [] + for char in text: + output.append(char) + while True: + escape_start = -1 + digits = "" + if ( + len(output) >= 6 + and output[-6] == "\\" + and output[-5] in ("u", "U") + and all(char in _ASCII_HEX_DIGITS for char in output[-4:]) + ): + escape_start = len(output) - 6 + digits = "".join(output[-4:]) + elif ( + len(output) >= 4 + and output[-4] == "\\" + and output[-3] in ("x", "X") + and all(char in _ASCII_HEX_DIGITS for char in output[-2:]) + ): + escape_start = len(output) - 4 + digits = "".join(output[-2:]) + elif ( + len(output) >= 3 + and output[-3] == "%" + and all(char in _ASCII_HEX_DIGITS for char in output[-2:]) + ): + escape_start = len(output) - 3 + digits = "".join(output[-2:]) + if escape_start < 0: + break + decoded = int(digits, 16) + if decoded > 0x7F: + break + del output[escape_start:] + output.append(chr(decoded)) + return "".join(output) + + +def _redact_sensitive_patterns(text: str) -> tuple[str, bool]: + """Redacts plain-text credential constructs without decoding the input.""" + + def _redact_key_value(match: re.Match[str]) -> str: + value = match.group("value") + if value == "[REDACTED]": + return match.group(0) + quote_char = value[:1] if value[:1] in ('"', "'") else "" + return f"{match.group('prefix')}{quote_char}[REDACTED]{quote_char}" + + def _redact_basic_credential(match: re.Match[str]) -> str: + encoded = match.group("value") + if len(encoded) > _MAX_BASIC_AUTH_TOKEN_CHARS: + return f"{match.group('prefix')}[REDACTED]" + try: + decoded = base64.b64decode(encoded, validate=True) + except (ValueError, TypeError): + return match.group(0) + if "=" not in encoded and b":" not in decoded: + return match.group(0) + return f"{match.group('prefix')}[REDACTED]" + + sanitized = _AUTH_HEADER_RE.sub( + lambda match: f"{match.group('prefix')}[REDACTED]", text + ) + sanitized = _BEARER_TOKEN_RE.sub( + lambda match: f"{match.group('prefix')}[REDACTED]", sanitized + ) + sanitized = _BASIC_TOKEN_RE.sub(_redact_basic_credential, sanitized) + sanitized = _SENSITIVE_TEXT_KEY_RE.sub(_redact_key_value, sanitized) + return sanitized, sanitized != text + + +def _contains_sensitive_text_marker(text: str) -> bool: + """Returns whether raw or encoded text contains a credential construct.""" + _, changed = _redact_sensitive_patterns(text) + if changed: + return True + canonical = _canonicalize_common_ascii_escapes(text) + if canonical == text: + return False + _, changed = _redact_sensitive_patterns(canonical) + return changed + + +def _sanitize_sensitive_text(text: str, max_len: int) -> tuple[str, bool]: + """Redacts bounded credential material embedded in diagnostic text. + + Unlike the generic attribute sanitizer, this preserves ordinary prose + exactly (including bracket-led log messages). Complete JSON/encoded JSON + still uses the existing structural redactor, while HTTP headers, bearer + tokens, query parameters, and key/value fragments embedded in prose are + replaced in place. Inputs too large to inspect safely fail closed when they + would otherwise be emitted without a configured length bound. + """ + if type(text) is not str: + text = str.__str__(text) + if len(text) > _MAX_JSON_INSPECT_CHARS: + if max_len == -1 or max_len > _MAX_JSON_INSPECT_CHARS: + return "[REDACTED_SENSITIVE_TEXT]", True + # Only this prefix can reach the row, so inspect exactly that bounded + # value below. Incomplete trailing escapes cannot reveal omitted bytes. + emitted = text[:max_len] + text = emitted + length_truncated = True + else: + length_truncated = False + + sanitized = text + stripped = _strip_bom_ws(text) + if stripped.startswith(("{", "[", '"')): + try: + json.loads(stripped) + except (TypeError, ValueError, RecursionError, MemoryError): + # Malformed diagnostic prose is handled by the explicit marker pass + # below; safe "[INFO] ..." messages must remain byte-identical. + pass + else: + structured, _ = _recursive_smart_truncate(text, -1) + if isinstance(structured, str): + sanitized = structured + else: + sanitized = json.dumps(structured) + + changed = sanitized != text + + sanitized, patterns_changed = _redact_sensitive_patterns(sanitized) + changed = changed or patterns_changed + + # Escaped/percent-encoded credential constructs cannot be safely rewritten + # in place without risking a partial source-to-canonical mapping. Fail the + # bounded diagnostic closed only when decoding actually exposes one; safe + # Windows paths and decoder errors remain byte-identical. + canonical = _canonicalize_common_ascii_escapes(sanitized) + if canonical != sanitized and _contains_sensitive_text_marker(canonical): + return "[REDACTED_SENSITIVE_TEXT]", True + + if max_len != -1 and len(sanitized) > max_len: + sanitized = sanitized[:max_len] + "...[TRUNCATED]" + length_truncated = True + return sanitized, changed or length_truncated + + # 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. +# 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 @@ -428,83 +646,357 @@ def _extract_tool_declarations( # 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. The remainder is +# 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 +# Hard ceiling on how many characters of a JSON-shaped string the +# synchronous sanitizer will materialize with json.loads, applied even when +# emitted content is unlimited (max_content_length=-1): a multi-megabyte +# encoded array otherwise burns unbounded event-loop time and memory before +# the node budget can apply (#6360 review round 8 P2-13). Container-capable +# values beyond it fail closed. +_MAX_JSON_INSPECT_CHARS = 4_000_000 + +# Keep deeply nested JSON behavior deterministic across Python versions. +# CPython <=3.13 rejects extreme nesting from json.loads with RecursionError, +# while 3.14's iterative decoder accepts it. A fixed ceiling preserves the +# fail-closed contract and bounds the materialized object graph everywhere. +_MAX_JSON_NESTING_DEPTH = 1_000 + + +def _strip_bom_ws(value: str) -> str: + """Strips BOMs and ALL Unicode whitespace from the start of a string. + + json.dumps round-trips arbitrary Unicode, so a decoded string layer can + hide a credential container behind U+00A0/U+2003-style whitespace that an + ASCII-only lstrip never removes (#6360 review round 8 P1-1). One linear + index scan: the earlier alternating lstrip loop re-sliced the remaining + suffix per whitespace/BOM pair, going quadratic on a legal prefix and + pinning the callback event loop (#6360 review round 9 P1-3). str.isspace + covers every Unicode whitespace; BOM (U+FEFF) is not whitespace, so both + are tested per character. + """ + i = 0 + n = len(value) + while i < n and (value[i].isspace() or value[i] == "\ufeff"): + i += 1 + return value[i:] if i else value + -def _json_nesting_exceeds(s: str, limit: int) -> bool: - """Reports whether JSON structural nesting in ``s`` exceeds ``limit``. +def _json_nesting_exceeds_limit(value: str) -> bool: + """Returns whether JSON structural nesting exceeds the fixed ceiling. - Scans bracket depth outside of string literals. ``json.loads``' own - recursion handling is interpreter-version dependent (CPython 3.14 parses - nesting that earlier versions reject with ``RecursionError``), so callers - bound the structural depth explicitly instead of relying on that error. + Brackets and braces inside strings are payload, not structure. This linear + preflight runs only after the existing character-size bound, before + ``json.loads`` can materialize a runtime-dependent deep object graph. """ depth = 0 in_string = False escaped = False - for ch in s: + for char in value: if in_string: if escaped: escaped = False - elif ch == "\\": + elif char == "\\": escaped = True - elif ch == '"': + elif char == '"': in_string = False continue - if ch == '"': + if char == '"': in_string = True - elif ch in "[{": + elif char in "[{": depth += 1 - if depth > limit: + if depth > _MAX_JSON_NESTING_DEPTH: return True - elif ch in "]}": + elif char in "]}": depth -= 1 return False +def _normalize_json_native( + obj: Any, + max_len: int, + depth: int = 0, + budget: Optional[list[int]] = None, +) -> tuple[Any, bool]: + """Coerces already-sanitized parser output to strictly JSON-native values. + + A nested hostile model can defer its failure PAST parser sanitization — + e.g. an attribute that returns an object whose __repr__ raises — so the + payload detonated later during Arrow serialization's json.dumps/str + fallback (#6360 review round 14 P1-2). Unlike the full sanitizer, this + does NOT re-run JSON-blob inspection on strings (parser output would + have its sentinels and bracketed prose corrupted), but it DOES enforce + the two policies parser output cannot be assumed to carry (#6360 review + round 15 P1-2/P1-3): sensitive-key/``temp:`` value redaction — a nested + model property can hand the parser a raw credential mapping — and the + configured length bound, because model fields like ``role`` are copied + verbatim and an unbounded string reopens the synchronous + json.dumps/Arrow work boundary. Returns ``(normalized, + replaced_anything)``. + """ + if budget is None: + budget = [_MAX_SANITIZE_NODES] + budget[0] -= 1 + if budget[0] < 0: + return "[SANITIZE_BUDGET_EXCEEDED]", True + if depth >= _MAX_SANITIZE_DEPTH: + return "[MAX_DEPTH_EXCEEDED]", True + try: + if isinstance(obj, str): + text = obj if type(obj) is str else str.__str__(obj) + if max_len != -1 and len(text) > max_len: + return text[:max_len] + "...[TRUNCATED]", True + return text, False + if obj is None or type(obj) in (int, float, bool): + return obj, False + if isinstance(obj, bool): + return bool(obj), False + if isinstance(obj, int): + return int(obj), False + if isinstance(obj, float): + return float(obj), False + if isinstance(obj, dict): + out_dict: dict[str, Any] = {} + replaced = False + bad_keys = 0 + + def collision_safe_key(k: str) -> str: + """Allocates a unique key while preserving first-writer order.""" + nonlocal bad_keys, replaced + if k not in out_dict: + return k + bad_keys += 1 + candidate = f"[KEY_COLLISION_{bad_keys}]{k}" + while candidate in out_dict: + bad_keys += 1 + candidate = f"[KEY_COLLISION_{bad_keys}]{k}" + replaced = True + return candidate + + for k, v in obj.items(): + if budget[0] <= 0: + budget_key = collision_safe_key("[SANITIZE_BUDGET_EXCEEDED]") + out_dict[budget_key] = "[SANITIZE_BUDGET_EXCEEDED]" + replaced = True + break + redact_value = False + if isinstance(k, str): + if type(k) is not str: + k = str.__str__(k) + k_lower = k.lower().replace("-", "_") + if k_lower in _SENSITIVE_KEYS or k_lower.startswith("temp:"): + redact_value = True + else: + bad_keys += 1 + k = f"[UNSUPPORTED_KEY_{bad_keys}]" + replaced = True + # Preserve every value after key normalization. The first writer + # keeps the plain key; later colliders receive a marker that is + # itself re-allocated around genuine marker-like user keys. + k = collision_safe_key(k) + if redact_value: + budget[0] -= 1 + out_dict[k] = "[REDACTED]" + continue + norm_v, v_replaced = _normalize_json_native( + v, max_len, depth + 1, budget + ) + replaced = replaced or v_replaced + out_dict[k] = norm_v + return out_dict, replaced + if isinstance(obj, (list, tuple)): + out_list: list[Any] = [] + replaced = False + for item in obj: + if budget[0] <= 0: + out_list.append("[SANITIZE_BUDGET_EXCEEDED]") + replaced = True + break + norm_item, item_replaced = _normalize_json_native( + item, max_len, depth + 1, budget + ) + replaced = replaced or item_replaced + out_list.append(norm_item) + return out_list, replaced + return "[UNSUPPORTED_OBJECT]", True + except Exception: + return "[UNSUPPORTED_OBJECT]", True + + +def _sanitize_free_text( + text: str, + seen: set[int], + depth: int, + max_len: int, + budget: Optional[list[int]], +) -> tuple[str, bool, bool]: + """Handles text that may embed a JSON document ANYWHERE, not just at + its start: raw multi-document suffixes and decoded string layers. + + A literal container token after prose cannot be verified (the text has + machine-encoded provenance, so a consumer may parse it) and fails + closed; a quoted fragment is walked through the full blob sanitizer + from the first quote (#6360 review rounds 12 P1-1 / 13 P1-2). A stray + escape in the prose gap could shift a consumer's parse boundaries and + also fails closed. Quote-free, container-free prose passes through. + """ + if "{" in text or "[" in text: + return "[UNPARSEABLE_JSON_BLOB]", True, True + quote_idx = text.find('"') + if quote_idx == -1: + return text, False, False + gap = text[:quote_idx] + if "\\" in gap: + return "[UNPARSEABLE_JSON_BLOB]", True, True + tail = text[quote_idx:] + t_sanitized, changed, truncated = _sanitize_json_blob( + tail, seen, depth + 1, max_len, budget + ) + if not changed: + return text, False, False + return gap + t_sanitized, True, truncated + + def _sanitize_json_blob( value: str, seen: set[int], depth: int = 0, max_len: int = -1, budget: Optional[list[int]] = None, -) -> tuple[str, bool]: +) -> tuple[str, bool, 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. Decode + 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)``; + arrays of credential objects included. Returns ``(value, changed, + truncated)`` where ``truncated`` reports content loss inside the blob + (depth/budget replacement discards payload, #6360 review round 6 P2-6); strings that do not parse, or that need no redaction, are returned unchanged (no cosmetic re-serialization). """ - stripped = value.lstrip("\ufeff \t\r\n") + # BOM/Unicode-whitespace-aware normalization must run before EVERY + # shape check — an ASCII-only lstrip let a decoded layer hide its + # container behind U+00A0/U+2003 (#6360 review round 8 P1-1). + stripped = _strip_bom_ws(value) + # The inspection ceiling applies even in unlimited mode (#6360 review + # round 8 P2-13). + inspect_limit = ( + min(max_len, _MAX_JSON_INSPECT_CHARS) + if max_len != -1 + else _MAX_JSON_INSPECT_CHARS + ) + if stripped.startswith('"'): + # A JSON-encoded STRING layer: json.dumps applied twice leaves the + # secret container quoted-and-escaped, which the container check + # below cannot see (#6360 review round 6 P1-3). Decode the layer and + # re-enter this sanitizer so a decoded string that turns out to be a + # container goes through the same redaction path. Layer recursion is + # bounded by the depth cap, and each decode strictly shrinks the + # string. + if depth >= _MAX_SANITIZE_DEPTH: + return "[UNPARSEABLE_JSON_BLOB]", True, True + if len(stripped) > inspect_limit: + # Decoding would allocate beyond the limit, so classify what will + # actually be EMITTED — the max_len prefix after the caller's raw + # truncation, or the ENTIRE value in unlimited mode, where scanning + # only the inspection window let a credential document sit just + # past it (#6360 review round 15 P1-1, refining rounds 7/8/14). Any + # container token or escape in the emitted text means the output + # could retain (escaped) credential material — fail closed. The + # scan is a bounded linear pass over text already in memory. + # Escape-free, container-free prose is left to the caller's length + # truncation (or emitted whole in unlimited mode). + emitted = stripped if max_len == -1 else stripped[:max_len] + if "{" in emitted or "[" in emitted or "\\" in emitted: + return "[UNPARSEABLE_JSON_BLOB]", True, True + return value, False, False + suffix = "" + try: + decoded = json.loads(stripped) + except (TypeError, RecursionError, MemoryError): + return "[UNPARSEABLE_JSON_BLOB]", True, True + except ValueError: + # Not one complete JSON document. A valid quoted credential layer + # followed by trailing garbage still lands here, and classifying it + # as prose republished the secret (#6360 review round 7 P1-1): + # bounded-decode the LEADING string layer instead. + try: + decoded, end = json.JSONDecoder().raw_decode(stripped) + except ValueError: + # No leading complete JSON document. If the post-quote prefix can + # still represent an ENCODED CONTAINER, the value is malformed + # credential JSON (e.g. an encoded layer missing its final quote) + # and must fail closed like the direct container path does for + # malformed JSON (#6360 review round 10 P1-1). Anything else is + # quoted prose. + if _strip_bom_ws(stripped[1:])[:1] in ("{", "[", "\\", '"'): + return "[UNPARSEABLE_JSON_BLOB]", True, True + return value, False, False + suffix = stripped[end:] + if not isinstance(decoded, str): + return value, False, False + stripped_suffix = _strip_bom_ws(suffix) + if "{" in suffix or "[" in suffix or stripped_suffix.startswith("\\"): + # An unverified raw suffix can smuggle a credential container past + # a harmless quoted prefix ('"note" {"access_token":...}') (#6360 + # review round 8 P1-2). A literal container-token scan is sound for + # RAW text that no later step decodes; a leading stray escape + # cannot be classified and fails closed. + return "[UNPARSEABLE_JSON_BLOB]", True, True + # ANY quoted fragment in the trailing text may be an encoded JSON + # document whose decoded content hides a container behind Unicode + # escapes — immediately ('"note" "\\u007b..."', round 11 P1-1) or + # after a stretch of prose ('"note" then "\\u007b..."', round 12 + # P1-1): walk it with the shared free-text handling. + s_sanitized, s_changed, s_truncated = _sanitize_free_text( + suffix, seen, depth, max_len, budget + ) + if s_changed and s_sanitized == "[UNPARSEABLE_JSON_BLOB]": + return "[UNPARSEABLE_JSON_BLOB]", True, True + inner = _strip_bom_ws(decoded) + if inner.startswith(("{", "[", '"')): + p_sanitized, p_changed, p_truncated = _sanitize_json_blob( + decoded, seen, depth + 1, max_len, budget + ) + else: + # The DECODED layer can hide a container after prose too + # ('"note \\u007b...access\\u005ftoken...\\u007d"' decodes to + # 'note {"access_token":...}') — the escapes are gone after this + # decode, so the same anywhere-in-text handling applies (#6360 + # review round 13 P1-2). + p_sanitized, p_changed, p_truncated = _sanitize_free_text( + decoded, seen, depth + 1, max_len, budget + ) + if p_changed and p_sanitized == "[UNPARSEABLE_JSON_BLOB]": + return "[UNPARSEABLE_JSON_BLOB]", True, True + if not p_changed and not s_changed: + return value, False, False + prefix_text = stripped[:end] if not p_changed else json.dumps(p_sanitized) + suffix_text = suffix if not s_changed else s_sanitized + return prefix_text + suffix_text, True, p_truncated or s_truncated if not stripped.startswith(("{", "[")): - return value, False + return value, False, 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. + # Enforce the inspection 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 + if len(stripped) > inspect_limit: + return "[UNPARSEABLE_JSON_BLOB]", True, True - # Fail closed on nesting too deep to fully inspect. Relying on json.loads - # to raise RecursionError is interpreter-version dependent (CPython 3.14 - # parses depths that earlier versions reject), so bound structural depth - # explicitly: a blob deeper than the sanitizer can traverse cannot be - # verified secret-free. - if _json_nesting_exceeds(stripped, _MAX_SANITIZE_DEPTH): - return "[UNPARSEABLE_JSON_BLOB]", True + if _json_nesting_exceeds_limit(stripped): + return "[UNPARSEABLE_JSON_BLOB]", True, 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. Track duplicates while parsing and always reserialize + # 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 @@ -520,30 +1012,34 @@ def _pairs_hook(pairs: list[tuple[str, Any]]) -> dict[str, Any]: try: parsed = json.loads(stripped, object_pairs_hook=_pairs_hook) if not isinstance(parsed, (dict, list)): - return value, False + return value, False, 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( + # The nested truncation flag must survive: a depth/budget sentinel + # inside the blob discards payload, and dropping the bit made the row + # claim completeness (#6360 review round 6 P2-6). + sanitized, nested_truncated = _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 + if sanitized == parsed and not saw_duplicate_key and not nested_truncated: + return value, False, False + return json.dumps(sanitized), True, nested_truncated 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), integers over the interpreter digit limit, or a blob too - # deep/large to inspect. None of these can be verified secret-free — and - # a raw-substring fallback is bypassable via JSON string escapes — so + # 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 + return "[UNPARSEABLE_JSON_BLOB]", True, 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. + 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}.") @@ -551,28 +1047,19 @@ def _require_count(name: str, value: Any, minimum: int) -> None: raise ValueError(f"{name} must be >= {minimum}, got {value}.") -def _require_finite( - name: str, value: Any, minimum: float, *, inclusive: bool = False -) -> float: - """Requires a finite real number at or above a minimum bound. - - The bound is exclusive by default (``value`` must be strictly greater than - ``minimum``); pass ``inclusive=True`` to also accept ``value == minimum``. - """ +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 inclusive: - if value < minimum: - raise ValueError(f"{name} must be >= {minimum}, got {value}.") - elif value <= minimum: - raise ValueError(f"{name} must be > {minimum}, got {value}.") + 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. + """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 @@ -601,16 +1088,24 @@ def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: _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. + # max_retries=0, initial_delay=0, max_delay=0 (#6360 review round 4 P2). initial_delay = _require_finite( - "retry_config.initial_delay", retry.initial_delay, 0, inclusive=True - ) - multiplier = _require_finite( - "retry_config.multiplier", retry.multiplier, 1, inclusive=True - ) - max_delay = _require_finite( - "retry_config.max_delay", retry.max_delay, 0, inclusive=True + "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" @@ -627,6 +1122,76 @@ def _validate_runtime_config(config: "BigQueryLoggerConfig") -> None: ) +class _SetupAbortedError(RuntimeError): + """Setup lost the lifecycle-generation race against shutdown(). + + Distinct from service failures so waiters and row owners can classify + the outcome without string matching (#6360 review round 8 P2-9). + """ + + +class _LoopStateAdmissionAbortedError(_SetupAbortedError): + """A retained writer is terminal for admission but still draining. + + This is a lifecycle outcome, not a setup/service failure: setup owners and + coalesced waiters must report ``aborted`` without poisoning retry backoff or + tearing down shared clients needed by the in-flight drain. + """ + + +class _ShutdownIncompleteError(RuntimeError): + """The owning shutdown() was cancelled before teardown completed. + + Coalesced callers must not treat this as success (#6360 review round + 10 P1-3); they retry ownership instead. + """ + + +def _base_str(safe_base: Any, obj: Any) -> str: + """Non-overridable base-class string conversion. + + ``safe_base.__str__(obj)`` bypasses any subclass ``__str__`` override — + a subclassed allowlisted scalar (application Enum, PurePath, ...) could + otherwise reopen the arbitrary-string leak (#6360 review round 8 P1-3). + """ + text = safe_base.__str__(obj) + return text if isinstance(text, str) else "[UNSUPPORTED_OBJECT]" + + +def _safe_getattr(obj: Any, name: str) -> Any: + """getattr that cannot be weaponized by payload-controlled properties. + + A bare ``hasattr``/``getattr`` probe evaluates properties, and hasattr + only swallows AttributeError — a property raising anything else escaped + the sanitizer entirely, dropping the telemetry row and leaking the + exception message (often containing the payload) into application logs + (#6360 review round 7 P1-3). Returns None on ANY failure. + """ + try: + return getattr(obj, name, None) + except Exception: + return None + + +# Stdlib scalar types whose str() form is canonical, side-effect free, and +# cannot embed attribute state beyond the value itself. Only these keep the +# stringify fallback; arbitrary objects' repr/str output is payload- +# controlled (e.g. SimpleNamespace(access_token=...) prints its attributes +# verbatim) and is no longer published (#6360 review round 7 P1-3). +_SAFE_STR_TYPES: tuple[type, ...] = ( + datetime, + date, + timedelta, + decimal.Decimal, + uuid.UUID, + pathlib.PurePath, + # NOTE: enum.Enum is handled FIRST in the dispatch (before the scalar + # branches), not here — value-backed members would otherwise never + # reach this fallback (#6360 review round 9 P1-2). + complex, +) + + def _recursive_smart_truncate( obj: Any, max_len: int, @@ -661,7 +1226,7 @@ def _recursive_smart_truncate( # 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 — unlike "[CIRCULAR_REFERENCE]", + # 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 @@ -670,21 +1235,55 @@ def _recursive_smart_truncate( if obj_id in seen: return "[CIRCULAR_REFERENCE]", False - # Track compound objects to detect cycles + # Converter attributes are fetched exactly once through _safe_getattr: + # hasattr() evaluates payload-controlled properties OUTSIDE any guard, + # and a property raising e.g. RuntimeError escaped the sanitizer, leaked + # its message into logs, and suppressed the whole row (#6360 review + # round 7 P1-3). + model_dump_fn = _safe_getattr(obj, "model_dump") + dict_fn = _safe_getattr(obj, "dict") + to_dict_fn = _safe_getattr(obj, "to_dict") + instance_dict = None + if not isinstance( + obj, (str, bytes, bytearray, int, float, bool, dict, list, tuple) + ): + attrs = _safe_getattr(obj, "__dict__") + if isinstance(attrs, dict): + instance_dict = attrs + + # Track compound objects to detect cycles. Plain objects traversed via + # their __dict__ count too: an unmarked self-reference (obj.self = obj) + # repeated the full attribute copy at every level until the depth cap + # (#6360 review round 8 P2-8). is_compound = ( 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") - or hasattr(obj, "to_dict") + or model_dump_fn is not None + or dict_fn is not None + or to_dict_fn is not None + or instance_dict is not None ) if is_compound: seen.add(obj_id) try: - if isinstance(obj, str): - obj, blob_replaced = _sanitize_json_blob( + if isinstance(obj, enum.Enum): + # BEFORE the scalar branches: a StrEnum / (str, Enum) / bytes-backed + # member is also an instance of its mixin type, so scalar dispatch + # normalized it with str.__str__ and published the underlying VALUE + # instead of the member name (#6360 review round 9 P1-2). + return _recursive_smart_truncate( + _base_str(enum.Enum, obj), max_len, seen, depth + 1, budget + ) + elif isinstance(obj, str): + if type(obj) is not str: + # str subclasses can override lstrip/startswith/lower to + # misreport their content while json.dumps still serializes the + # real underlying value (#6360 review round 7 P1-3): normalize to + # the exact built-in string before any inspection. + obj = str.__str__(obj) + obj, blob_replaced, blob_truncated = _sanitize_json_blob( obj, seen, depth, max_len, budget ) if blob_replaced and obj == "[UNPARSEABLE_JSON_BLOB]": @@ -692,32 +1291,98 @@ def _recursive_smart_truncate( return obj, True if max_len != -1 and len(obj) > max_len: return obj[:max_len] + "...[TRUNCATED]", True - return obj, False + return obj, blob_truncated elif isinstance(obj, (bytes, bytearray)): # Credential JSON frequently travels as bytes; stringifying it in - # the fallback bypassed blob redaction. + # the fallback bypassed blob redaction (#6360 review round 5 P1-4). try: decoded = bytes(obj).decode("utf-8") except UnicodeDecodeError: - return "[BINARY_DATA]", False + # The whole byte payload is discarded, so the row must report + # content loss (#6360 review round 7 P2-9). + return "[BINARY_DATA]", True 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. - # Always emits a plain sanitized dict. + # 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. new_dict = {} + unsupported_keys = 0 for k, v in obj.items(): + # Stop iterating once the work budget is exhausted: recursing on + # every remaining entry still did O(input) work and produced + # O(input) sentinel output, and directly-redacted entries consumed + # no budget at all, so a wide "temp:" mapping bypassed the bound + # entirely (#6360 review round 6 P2-5). One remainder sentinel + # stands in for everything dropped. + if budget[0] <= 0: + new_dict["[SANITIZE_BUDGET_EXCEEDED]"] = "[SANITIZE_BUDGET_EXCEEDED]" + truncated_any = True + break + redact_value = False if isinstance(k, str): - k_lower = k.lower() + if type(k) is not str: + # Same normalization as string values: a str-subclass key can + # misreport itself to the redaction check below (#6360 review + # round 7 P1-3). + k = str.__str__(k) + k_lower = k.lower().replace("-", "_") if k_lower in _SENSITIVE_KEYS or k_lower.startswith("temp:"): - new_dict[k] = "[REDACTED]" - continue + redact_value = True + elif k is None or isinstance(k, (int, float, bool)): + # JSON stringifies all object keys, so 1 and "1" (or True and + # "true", None and "null") silently collapse into duplicate + # members and one value is lost downstream (#6360 review round + # 8 P2-10). Normalize to the exact JSON key form BEFORE + # insertion; collisions are handled below. + if k is True: + k = "true" + elif k is False: + k = "false" + elif k is None: + k = "null" + else: + k = json.dumps(k) + else: + # json.dumps rejects non-scalar keys even with default=str (it + # applies to values only), so one such key raised TypeError at + # serialization and silently dropped the whole telemetry row + # (#6360 review round 7 P2-8). The key's own repr can embed + # secrets, so it is not published either: fail the KEY closed + # and keep the sanitized value. + unsupported_keys += 1 + budget[0] -= 1 + k = ( + "[UNSUPPORTED_KEY]" + if unsupported_keys == 1 + else f"[UNSUPPORTED_KEY_{unsupported_keys}]" + ) + truncated_any = True + if k in new_dict: + # Deterministic fail-closed collision policy (#6360 review + # round 8 P2-10): the first writer keeps the plain key; later + # colliders keep their (sanitized) value under an explicit + # marker instead of silently overwriting or being dropped at + # JSON parse time. The marker is re-allocated until unique — a + # single fixed marker could alias (and overwrite) a legitimate + # user key already named "[KEY_COLLISION_n]..." (#6360 review + # round 9 P2-9). + unsupported_keys += 1 + candidate = f"[KEY_COLLISION_{unsupported_keys}]{k}" + while candidate in new_dict: + unsupported_keys += 1 + candidate = f"[KEY_COLLISION_{unsupported_keys}]{k}" + k = candidate + truncated_any = True + + if redact_value: + budget[0] -= 1 + new_dict[k] = "[REDACTED]" + continue val, trunc = _recursive_smart_truncate( v, max_len, seen, depth + 1, budget ) @@ -726,10 +1391,29 @@ def _recursive_smart_truncate( new_dict[k] = val return new_dict, truncated_any elif isinstance(obj, (list, tuple)): + fields = _safe_getattr(obj, "_fields") if isinstance(obj, tuple) else None + if ( + isinstance(fields, tuple) + and len(fields) == len(obj) + and all(isinstance(f, str) for f in fields) + ): + # Namedtuples carry field NAMES (e.g. access_token) that the + # plain-list normalization below discarded, so the sensitive + # value survived positionally (#6360 review round 7 P1-2). + # Rebuild as a mapping from the class metadata (zip, not the + # overridable _asdict()) so key redaction runs. + return _recursive_smart_truncate( + dict(zip(fields, obj)), max_len, seen, depth + 1, budget + ) truncated_any = False new_list = [] # Explicit loop to handle flag propagation for i in obj: + # Same bound as the mapping loop (#6360 review round 6 P2-5). + if budget[0] <= 0: + new_list.append("[SANITIZE_BUDGET_EXCEEDED]") + truncated_any = True + break val, trunc = _recursive_smart_truncate( i, max_len, seen, depth + 1, budget ) @@ -740,7 +1424,7 @@ def _recursive_smart_truncate( return type(obj)(new_list), truncated_any # Tuple/list subclasses (e.g. namedtuples) may require positional # constructor fields; reconstructing raised TypeError and the safe - # callback dropped the whole row. JSON + # 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 @@ -750,59 +1434,118 @@ def _recursive_smart_truncate( return _recursive_smart_truncate( as_dict, max_len, seen, depth + 1, budget ) - elif hasattr(obj, "model_dump") and callable(obj.model_dump): + elif model_dump_fn is not None and callable(model_dump_fn): # Pydantic v2. Only recurse if the conversion made PROGRESS toward a - # JSON-native container: Mock-like objects answer every duck-typed + # JSON-native value: 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. + # at the fallback. Progress includes SCALARS: a RootModel[str] dumps + # to a plain string that may itself be a credential blob, and + # falling through to the generic fallback bypassed blob redaction + # (#6360 review round 6 P1-3). try: - dumped = obj.model_dump() - if isinstance(dumped, (collections.abc.Mapping, list)): + dumped = model_dump_fn() + if isinstance( + dumped, + (collections.abc.Mapping, list, tuple, str, bytes, bytearray), + ): return _recursive_smart_truncate( dumped, max_len, seen, depth + 1, budget ) + if dumped is None or isinstance(dumped, (int, float, bool)): + return dumped, False except Exception: pass - elif hasattr(obj, "dict") and callable(obj.dict): - # Pydantic v1 (same progress requirement as above). + elif dict_fn is not None and callable(dict_fn): + # Pydantic v1 (same progress requirement as above, scalars included). try: - dumped = obj.dict() - if isinstance(dumped, (collections.abc.Mapping, list)): + dumped = dict_fn() + if isinstance( + dumped, + (collections.abc.Mapping, list, tuple, str, bytes, bytearray), + ): return _recursive_smart_truncate( dumped, max_len, seen, depth + 1, budget ) + if dumped is None or isinstance(dumped, (int, float, bool)): + return dumped, False except Exception: pass - elif hasattr(obj, "to_dict") and callable(obj.to_dict): - # Common pattern for custom objects (same progress requirement). + elif to_dict_fn is not None and callable(to_dict_fn): + # Common pattern for custom objects (same progress requirement, + # scalars included). try: - dumped = obj.to_dict() - if isinstance(dumped, (collections.abc.Mapping, list)): + dumped = to_dict_fn() + if isinstance( + dumped, + (collections.abc.Mapping, list, tuple, str, bytes, bytearray), + ): return _recursive_smart_truncate( dumped, max_len, seen, depth + 1, budget ) + if dumped is None or isinstance(dumped, (int, float, bool)): + return dumped, False 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, then RE-ENTER the - # string sanitizer — an object whose __str__ returns credential JSON - # bypassed blob redaction otherwise. - # 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 + # Fallback for unknown types. Arbitrary str()/repr() output is + # payload-controlled and prints attribute values verbatim (e.g. + # SimpleNamespace(access_token=...)), so it is no longer published + # (#6360 review round 7 P1-3; supersedes the round-5 stringify + # fallback). Known-safe stdlib scalars keep their canonical string + # form via NON-OVERRIDABLE base-class conversions — a subclass + # (application Enum, PurePath, ...) can override __str__ to return + # its value, reopening the arbitrary-string leak through a + # polymorphic str(obj) (#6360 review round 8 P1-3). Objects exposing + # instance state keep a structurally sanitized public view of their + # __dict__, collected through a budget-checked loop (a million- + # attribute object was fully copied before the budget applied, #6360 + # review round 8 P2-8); everything else fails closed to a sentinel. + # Sentinel replacement discards content, so it reports truncation. + for safe_base in _SAFE_STR_TYPES: + if isinstance(obj, safe_base): + # _base_str bypasses any subclass override; propagate the + # truncation flag — dropping it made an over-limit safe scalar + # claim completeness (#6360 review round 8 P2-11). + return _recursive_smart_truncate( + _base_str(safe_base, obj), max_len, seen, depth + 1, budget + ) + if instance_dict is not None: + public_attrs = {} + overflow = False + for k, v in instance_dict.items(): + if budget[0] <= 0: + overflow = True + break + # Every entry costs budget, filtered or not: skipping private + # attributes for free left the collection loop O(input). + budget[0] -= 1 + if isinstance(k, str) and not k.startswith("_"): + public_attrs[k] = v + if public_attrs or overflow: + sanitized_attrs, attrs_truncated = _recursive_smart_truncate( + public_attrs, max_len, seen, depth + 1, budget + ) + if overflow and isinstance(sanitized_attrs, dict): + sanitized_attrs["[SANITIZE_BUDGET_EXCEEDED]"] = ( + "[SANITIZE_BUDGET_EXCEEDED]" + ) + attrs_truncated = True + return sanitized_attrs, attrs_truncated + return "[UNSUPPORTED_OBJECT]", True + except Exception: + # Fail-closed protocol boundary (#6360 review round 8 P1-4): a + # hostile container protocol (a Mapping whose items() raises, + # sequence iteration or dataclass field access raising) must neither + # escape the sanitizer — the safe callback would log the payload- + # controlled message and drop the whole row — nor be logged here. + return "[UNSUPPORTED_OBJECT]", True finally: if is_compound: - seen.remove(obj_id) + seen.discard(obj_id) # --- PyArrow Helper Functions --- @@ -1483,6 +2226,11 @@ def __init__( self._queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue( maxsize=queue_max_size ) + self._queue_max_size = queue_max_size + # Outstanding shutdown sentinels currently in the queue: lets + # cancellation account for remaining rows in O(1) via qsize() instead + # of a synchronous full drain (#6360 review round 8 P2-12). + self._sentinel_count = 0 self._batch_processor_task: Optional[asyncio.Task[None]] = None self._shutdown = False @@ -1498,6 +2246,7 @@ def __init__( "non_retryable": 0, "unexpected_error": 0, "shutdown_timeout": 0, + "shutdown_cancelled": 0, } async def flush(self) -> None: @@ -1539,6 +2288,9 @@ def get_drop_stats(self) -> dict[str, int]: ``non_retryable``: BigQuery returned a non-retryable error (e.g. a schema mismatch). ``unexpected_error``: an unexpected exception aborted the write. + ``shutdown_timeout``: rows still queued when shutdown timed out. + ``shutdown_cancelled``: rows still queued when shutdown was + cancelled from outside (e.g. a host close timeout). Returns: A copy of the per-reason drop counters. @@ -1638,6 +2390,7 @@ async def _batch_writer(self) -> None: ) if first_item is _SHUTDOWN_SENTINEL: + self._sentinel_count = max(0, self._sentinel_count - 1) self._queue.task_done() continue @@ -1647,6 +2400,7 @@ async def _batch_writer(self) -> None: try: item = self._queue.get_nowait() if item is _SHUTDOWN_SENTINEL: + self._sentinel_count = max(0, self._sentinel_count - 1) self._queue.task_done() continue batch.append(item) @@ -1665,7 +2419,7 @@ async def _batch_writer(self) -> None: continue except asyncio.CancelledError: # Cancelled (e.g. by the shutdown timeout): the in-flight batch is - # lost — count it — then exit the + # 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) @@ -1778,7 +2532,10 @@ async def perform_write() -> None: if row_errors: for row_error in row_errors: logger.error("Row error details: %s", row_error) - logger.error("Row content causing error: %s", rows) + logger.error( + "%d row(s) dropped due to a non-retryable BigQuery error.", + len(rows), + ) self._dropped["non_retryable"] += len(rows) return return @@ -1827,6 +2584,24 @@ async def perform_write() -> None: ) return + def _drain_queue_and_count(self, reason: str) -> int: + """Counts and discards everything still queued (loss accounting).""" + drained = 0 + try: + while True: + item = self._queue.get_nowait() + if item is not _SHUTDOWN_SENTINEL: + drained += 1 + else: + self._sentinel_count = max(0, self._sentinel_count - 1) + self._queue.task_done() + except asyncio.QueueEmpty: + pass + if drained: + self._dropped[reason] += drained + logger.warning("%d queued row(s) dropped (%s).", drained, reason) + return drained + async def shutdown(self, timeout: float = 5.0) -> None: """Shuts down the BatchProcessor, draining the queue. @@ -1839,38 +2614,67 @@ async def shutdown(self, timeout: float = 5.0) -> None: # Signal the writer to wake up and check shutdown status try: self._queue.put_nowait(_SHUTDOWN_SENTINEL) + self._sentinel_count += 1 except asyncio.QueueFull: # If queue is full, the writer is active and will check _shutdown soon pass if self._batch_processor_task: + if self._batch_processor_task.done(): + # A previous shutdown attempt already terminated the worker — + # possibly by external cancellation. Re-awaiting the task would + # re-raise its historical CancelledError on EVERY retry, so no + # later close could ever finish cleanup (#6360 review round 7 + # P1-4). Treat the terminal worker as final and account for + # whatever is still queued. + self._drain_queue_and_count( + "shutdown_cancelled" + if self._batch_processor_task.cancelled() + else "shutdown_timeout" + ) + return try: await asyncio.wait_for(self._batch_processor_task, timeout=timeout) except asyncio.TimeoutError: logger.warning("BatchProcessor shutdown timed out, cancelling worker.") self._batch_processor_task.cancel() - try: - # Wait for the task to acknowledge cancellation - await self._batch_processor_task - except asyncio.CancelledError: - pass + # Convert the WORKER's expected CancelledError into a gather result, + # then shield that acknowledgement owner. If the HOST cancels this + # shutdown await, shield raises CancelledError unambiguously while the + # gather continues to own/retrieve the worker result. This works on + # Python 3.10 (where Task.cancelling() does not exist) and preserves + # the external-cancellation distinction on newer runtimes. + await asyncio.shield( + asyncio.gather( + self._batch_processor_task, + return_exceptions=True, + ) + ) # Rows still queued after the timeout are lost: count them so the - # loss is observable instead of silent. + # 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 + self._drain_queue_and_count("shutdown_timeout") + except asyncio.CancelledError: + # EXTERNAL cancellation (e.g. PluginManager's close timeout): + # wait_for has already cancelled the worker. Account for queued + # rows now — a retry may never come — and preserve the caller's + # cancellation (#6360 review round 7 P1-4). The accounting is + # O(1): this handler runs INSIDE the caller's cancellation window + # (asyncio.timeout waits for cleanup), so a synchronous full + # drain of an unbounded queue extended host-close latency + # linearly with queue depth (#6360 review round 8 P2-12). qsize + # minus outstanding sentinels counts the rows; storage is + # released by swapping in a fresh queue, reclaimed by GC off the + # cancellation-critical path. + remaining = max(0, self._queue.qsize() - self._sentinel_count) + if remaining: + self._dropped["shutdown_cancelled"] += remaining logger.warning( - "%d queued row(s) dropped by shutdown timeout.", drained + "%d queued row(s) dropped (shutdown_cancelled).", remaining ) + self._queue = asyncio.Queue(maxsize=self._queue_max_size) + self._sentinel_count = 0 + raise except Exception as e: logger.error("Error during BatchProcessor shutdown: %s", e) @@ -1896,7 +2700,8 @@ async def close(self) -> None: except asyncio.CancelledError: pass # Same loss accounting as shutdown(): rows still queued after the - # timeout are counted, not silently discarded. The cancelled worker counts its own in-flight batch. + # 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: @@ -1962,7 +2767,8 @@ def _upload_sync( # 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. + # 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 ) @@ -1995,6 +2801,186 @@ def _truncate(self, text: str) -> tuple[str, bool]: ) return text, False + @staticmethod + def _sanitize_raw_text(text: str) -> tuple[str, bool]: + """Sanitizes caller-provided text exactly once before it is stored. + + Known formatter/redaction sentinels are already safe and must not be + reinterpreted as malformed JSON merely because they are bracketed. All + other raw text goes through the same fail-closed credential sanitizer used + for structured attributes. Length truncation remains a separate pass so + the sanitized value is also the value sent to GCS. + """ + if text in (_FORMATTER_FAILED_SENTINEL, "[REDACTED]"): + return text, False + sanitized, content_lost = _recursive_smart_truncate(text, -1) + if sanitized == "[UNPARSEABLE_JSON_BLOB]": + # Raw content is user-facing prose, not an opaque attributes blob. + # Reusing the attributes sanitizer made every invalid bracket-led + # message ("[INFO]", Markdown links, "{not json}") disappear. Restore + # only bounded prose with no raw or encoded credential construct; + # malformed credential documents remain fail-closed while ordinary + # Windows paths and decoder diagnostics stay byte-identical. + stripped = _strip_bom_ws(text) + if ( + len(stripped) <= _MAX_JSON_INSPECT_CHARS + and stripped.startswith(("[", "{")) + and not _contains_sensitive_text_marker(stripped) + ): + try: + json.loads(stripped) + except (ValueError, RecursionError, MemoryError): + return text, False + if not isinstance(sanitized, str): + return "[UNSUPPORTED_OBJECT]", True + return sanitized, content_lost + + def _sanitize_and_truncate(self, text: str) -> tuple[str, bool]: + sanitized, content_lost = self._sanitize_raw_text(text) + truncated_text, length_truncated = self._truncate(sanitized) + return truncated_text, content_lost or length_truncated + + def _sanitize_external_uri(self, uri: str) -> tuple[str, bool]: + """Redacts signed/query credentials while preserving a URI's location.""" + if not isinstance(uri, str): + return "[REDACTED_SENSITIVE_URI]", True + if type(uri) is not str: + uri = str.__str__(uri) + if len(uri) > _MAX_JSON_INSPECT_CHARS: + return "[REDACTED_SENSITIVE_URI]", True + try: + parsed = urlsplit(uri) + if parsed.username is not None or parsed.password is not None: + # Userinfo is a credential-bearing URI surface by definition. Do not + # try to retain a username while guessing whether it is sensitive. + return "[REDACTED_SENSITIVE_URI]", True + query = parse_qsl(parsed.query, keep_blank_values=True) + except ValueError: + return "[REDACTED_SENSITIVE_URI]", True + + changed = False + path_segments = parsed.path.split("/") + redact_next_path_segment = False + for index, segment in enumerate(path_segments): + if not segment: + continue + canonical_segment = _canonicalize_common_ascii_escapes(segment) + if redact_next_path_segment: + path_segments[index] = quote("[REDACTED]", safe="") + changed = True + redact_next_path_segment = False + continue + if _is_sensitive_text_key(canonical_segment): + path_segments[index] = quote("[REDACTED]", safe="") + changed = True + redact_next_path_segment = True + continue + safe_segment, segment_changed = _sanitize_sensitive_text(segment, -1) + if segment_changed: + path_segments[index] = quote(safe_segment, safe="") + changed = True + + safe_query: list[tuple[str, str]] = [] + for key, value in query: + if _is_sensitive_text_key(key): + safe_query.append((key, "[REDACTED]")) + changed = True + continue + safe_key, key_changed = _sanitize_sensitive_text(key, -1) + safe_value, value_changed = _sanitize_sensitive_text(value, -1) + safe_query.append((safe_key, safe_value)) + changed = changed or key_changed or value_changed + + safe_fragment, fragment_changed = _sanitize_sensitive_text( + parsed.fragment, -1 + ) + changed = changed or fragment_changed + safe_uri = urlunsplit(( + parsed.scheme, + parsed.netloc, + "/".join(path_segments), + urlencode(safe_query), + safe_fragment, + )) + safe_uri, uri_truncated = self._truncate(safe_uri) + return safe_uri, changed or uri_truncated + + def _serialize_part_model(self, value: Any) -> tuple[dict[str, Any], bool]: + """Returns bounded JSON-native fields for a supported structured part.""" + dumped = value.model_dump(exclude_none=True, mode="json") + sanitized, content_lost = _recursive_smart_truncate(dumped, self.max_length) + if not isinstance(sanitized, dict): + return {"value": "[UNSUPPORTED_OBJECT]"}, True + + budget = [_MAX_SANITIZE_NODES] + + def _sanitize_strings(obj: Any, depth: int = 0) -> tuple[Any, bool]: + budget[0] -= 1 + if budget[0] < 0 or depth >= _MAX_SANITIZE_DEPTH: + return "[SANITIZE_BUDGET_EXCEEDED]", True + if isinstance(obj, str): + return _sanitize_sensitive_text(obj, self.max_length) + if isinstance(obj, dict): + out: dict[str, Any] = {} + replaced = False + collision_count = 0 + + def _collision_safe_key(key: str) -> str: + nonlocal collision_count, replaced + if key not in out: + return key + collision_count += 1 + candidate = f"[KEY_COLLISION_{collision_count}]{key}" + while candidate in out: + collision_count += 1 + candidate = f"[KEY_COLLISION_{collision_count}]{key}" + replaced = True + return candidate + + for key, item in obj.items(): + if budget[0] <= 0: + budget_key = _collision_safe_key("[SANITIZE_BUDGET_EXCEEDED]") + out[budget_key] = "[SANITIZE_BUDGET_EXCEEDED]" + return out, True + redact_item = False + if not isinstance(key, str): + safe_key = "[UNSUPPORTED_KEY]" + key_replaced = True + else: + if type(key) is not str: + key = str.__str__(key) + canonical_key = _canonicalize_common_ascii_escapes(key) + redact_item = _is_sensitive_text_key(canonical_key) + safe_key, key_replaced = _sanitize_sensitive_text( + key, self.max_length + ) + safe_key = _collision_safe_key(safe_key) + if redact_item: + safe_item = "[REDACTED]" + item_replaced = item != "[REDACTED]" + else: + safe_item, item_replaced = _sanitize_strings(item, depth + 1) + out[safe_key] = safe_item + replaced = replaced or key_replaced or item_replaced + return out, replaced + if isinstance(obj, list): + out_list = [] + replaced = False + for item in obj: + if budget[0] <= 0: + out_list.append("[SANITIZE_BUDGET_EXCEEDED]") + return out_list, True + safe_item, item_replaced = _sanitize_strings(item, depth + 1) + out_list.append(safe_item) + replaced = replaced or item_replaced + return out_list, replaced + return obj, False + + sanitized, text_content_lost = _sanitize_strings(sanitized) + if not isinstance(sanitized, dict): + return {"value": "[UNSUPPORTED_OBJECT]"}, True + return sanitized, content_lost or text_content_lost + async def _parse_content_object( self, content: types.Content | types.Part, @@ -2008,13 +2994,14 @@ async def _parse_content_object( ``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. They fall back to the + 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. + 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 @@ -2038,7 +3025,12 @@ async def _parse_content_object( # CASE A: It is already a URI (e.g. from user input) if hasattr(part, "file_data") and part.file_data: part_data["storage_mode"] = "EXTERNAL_URI" - part_data["uri"] = part.file_data.file_uri + safe_uri, uri_content_lost = self._sanitize_external_uri( + part.file_data.file_uri + ) + part_data["uri"] = safe_uri + if uri_content_lost: + is_truncated = True part_data["mime_type"] = part.file_data.mime_type # CASE B: It is Binary/Inline Data (Image/Blob) @@ -2074,8 +3066,11 @@ async def _parse_content_object( # CASE C: Text elif hasattr(part, "text") and part.text: - char_len = len(part.text) - byte_len = len(part.text.encode("utf-8")) + safe_text, sanitized_content_lost = self._sanitize_raw_text(part.text) + if sanitized_content_lost: + is_truncated = True + char_len = len(safe_text) + byte_len = len(safe_text.encode("utf-8")) # Decide whether to offload using each limit in its own # unit. inline_text_limit is a byte-based storage guard; @@ -2093,7 +3088,7 @@ async def _parse_content_object( ) try: uri = await self.offloader.upload_content( - part.text, "text/plain", path + safe_text, "text/plain", path ) part_data["storage_mode"] = "GCS_REFERENCE" part_data["uri"] = uri @@ -2107,17 +3102,17 @@ async def _parse_content_object( } part_data["object_ref"] = object_ref part_data["mime_type"] = "text/plain" - part_data["text"] = part.text[:200] + "... [OFFLOADED]" + part_data["text"] = safe_text[:200] + "... [OFFLOADED]" except Exception as e: logger.warning("Failed to offload text to GCS: %s", e) - clean_text, truncated = self._truncate(part.text) + clean_text, truncated = self._truncate(safe_text) if truncated: is_truncated = True part_data["text"] = clean_text summary_text.append(clean_text) else: # Text is small or no offloader, keep inline - clean_text, truncated = self._truncate(part.text) + clean_text, truncated = self._truncate(safe_text) if truncated: is_truncated = True part_data["text"] = clean_text @@ -2130,6 +3125,66 @@ async def _parse_content_object( {"function_name": part.function_call.name} ) + elif hasattr(part, "function_response") and part.function_response: + response, response_lost = self._serialize_part_model( + part.function_response + ) + if response_lost: + is_truncated = True + name = response.get("name") or "unknown" + if not isinstance(name, str): + name = "[UNSUPPORTED_OBJECT]" + is_truncated = True + response_summary = f"Function response: {name}" + part_data["mime_type"] = "application/json" + part_data["text"] = response_summary + part_data["part_attributes"] = json.dumps( + {"function_response": response} + ) + summary_text.append(response_summary) + + elif hasattr(part, "executable_code") and part.executable_code: + executable, code_lost = self._serialize_part_model(part.executable_code) + if code_lost: + is_truncated = True + language = executable.get("language") or "unknown" + if not isinstance(language, str): + language = "[UNSUPPORTED_OBJECT]" + is_truncated = True + code = executable.get("code") or "" + if not isinstance(code, str): + code = "[UNSUPPORTED_OBJECT]" + is_truncated = True + part_data["mime_type"] = "text/plain" + part_data["text"] = code + part_data["part_attributes"] = json.dumps({ + "executable_code": executable, + }) + summary_text.append(f"Executable code ({language}): {code}") + + elif ( + hasattr(part, "code_execution_result") and part.code_execution_result + ): + result, result_lost = self._serialize_part_model( + part.code_execution_result + ) + if result_lost: + is_truncated = True + outcome = result.get("outcome") or "unknown" + if not isinstance(outcome, str): + outcome = "[UNSUPPORTED_OBJECT]" + is_truncated = True + output = result.get("output") or "" + if not isinstance(output, str): + output = "[UNSUPPORTED_OBJECT]" + is_truncated = True + part_data["mime_type"] = "text/plain" + part_data["text"] = output + part_data["part_attributes"] = json.dumps({ + "code_execution_result": result, + }) + summary_text.append(f"Code execution result ({outcome}): {output}") + content_parts.append(part_data) summary_str, truncated = self._truncate(" | ".join(summary_text)) @@ -2150,20 +3205,21 @@ async def parse( ``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. The instance - fields remain only as a backward-compatible default. + 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. + # 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 def process_text(t: str) -> tuple[str, bool]: - return self._truncate(t) + return self._sanitize_and_truncate(t) if isinstance(content, LlmRequest): # Handle Prompt @@ -2175,6 +3231,10 @@ def process_text(t: str) -> tuple[str, bool]: ) for content_idx, c in enumerate(contents): role = getattr(c, "role", "unknown") + if isinstance(role, str): + role, role_truncated = process_text(role) + if role_truncated: + is_truncated = True summary, parts, trunc = await self._parse_content_object( c, trace_id=trace_id, @@ -2879,19 +3939,29 @@ def __init__( self._startup_error: Optional[Exception] = None self._setup_failures = 0 self._setup_retry_at = 0.0 - # Plugin-level loss accounting: counts drops that happen + # 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] = {} + # Guards every read-modify-write and snapshot of _local_drop_counts: + # events run on loops in different threads, and the unlocked + # increment/fold underreported losses under contention (#6360 review + # round 7 P2-7). + self._drop_counts_guard = threading.Lock() self._is_shutting_down = False # 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 + # Concurrent shutdown callers coalesce on the active owner's + # completion future instead of returning before teardown finished + # (#6360 review round 9 P1-5). + self._shutdown_future: Optional["ConcurrentFuture[None]"] = None # Lifecycle generation: shutdown() bumps it so an in-flight setup that - # completes afterwards cannot resurrect _started. + # completes afterwards cannot resurrect _started (#6360 round 5 P1-2). self._generation = 0 - # Guards ownership changes of _loop_state_by_loop: unsynchronized iteration raced concurrent insertion. + # 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 @@ -2905,25 +3975,9 @@ def __init__( self._init_pid = os.getpid() _LIVE_PLUGINS.add(self) - def _count_unwritten_queued_rows(self, state: _LoopState) -> int: - """Counts rows still queued on a processor whose loop is gone. - - Uses ``qsize()`` rather than ``get_nowait()``: on a queue bound to a - closed event loop, ``get_nowait()`` can raise "Event loop is closed" - while waking blocked putters, whereas ``qsize()`` only reads the queue - length and never touches the loop. These rows can never be written, so - they are counted as lost. ``qsize()`` may overcount by one if a shutdown - sentinel is still queued; that is acceptable for best-effort loss stats - and avoids depending on the private ``_queue`` attribute. - """ - queue = getattr(state.batch_processor, "_queue", None) - if not isinstance(queue, asyncio.Queue): - return 0 - return queue.qsize() - def _cleanup_stale_loop_states(self) -> None: """Removes entries for event loops that have been closed.""" - # Snapshot under the guard: iterating the + # 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. @@ -2931,11 +3985,36 @@ def _cleanup_stale_loop_states(self) -> None: candidates = list(self._loop_state_by_loop) stale = [loop for loop in candidates if loop.is_closed()] for loop in stale: - # Atomic claim: exactly one concurrent - # cleanup folds a given processor's counters — read-fold-delete - # raced, double-counting and raising KeyError. - with self._loop_states_guard: + # Atomic claim AND fold (#6360 review round 3 P2, round 8 P2-7): + # exactly one concurrent cleanup folds a given processor's counters, + # and the pop and the fold happen in one guarded transition so + # get_drop_stats() never observes the state as neither live nor + # folded (or as both). + stale_rows = 0 + with self._loop_states_guard, self._drop_counts_guard: state = self._loop_state_by_loop.pop(loop, None) + if state is not None: + 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). O(1) accounting (qsize minus tracked sentinels, + # #6360 review round 9 P2-10), INSIDE the single-winner claim: + # counting after the claim let a shutdown holding an earlier + # snapshot count the same queue a second time (#6360 review + # round 10 P2-5). + queue = getattr(state.batch_processor, "_queue", None) + sentinels = getattr(state.batch_processor, "_sentinel_count", 0) + if isinstance(queue, asyncio.Queue): + if not isinstance(sentinels, int): + sentinels = 0 + stale_rows = max(0, queue.qsize() - sentinels) + if stale_rows: + self._local_drop_counts["stale_loop"] = ( + self._local_drop_counts.get("stale_loop", 0) + stale_rows + ) if state is None: continue logger.warning( @@ -2943,19 +4022,7 @@ def _cleanup_stale_loop_states(self) -> None: loop, id(loop), ) - # Preserve the dead processor's loss accounting before discarding it, - # 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. - stale_rows = self._count_unwritten_queued_rows(state) 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) ) @@ -3047,9 +4114,45 @@ def _format_content_safely( logger.warning("Content formatter failed: %s", e) return "[FORMATTING FAILED]", False - async def _get_loop_state(self) -> _LoopState: + async def _close_write_transport(self, write_client: Any) -> None: + """Best-effort bounded close for a BigQuery write-client transport.""" + transport = getattr(write_client, "transport", None) + close_fn = getattr(transport, "close", None) + if close_fn is None: + return + try: + if asyncio.iscoroutinefunction(close_fn): + await asyncio.wait_for(close_fn(), timeout=self.config.shutdown_timeout) + else: + loop = asyncio.get_running_loop() + result = await asyncio.wait_for( + loop.run_in_executor(None, close_fn), + timeout=self.config.shutdown_timeout, + ) + if isinstance(result, collections.abc.Awaitable): + await asyncio.wait_for(result, timeout=self.config.shutdown_timeout) + except asyncio.CancelledError: + raise + except Exception: + logger.warning("Could not close a detached BigQuery write transport.") + + async def _close_detached_loop_transport(self, state: _LoopState) -> None: + """Best-effort bounded close for a terminal loop state's transport.""" + await self._close_write_transport(state.write_client) + + async def _get_loop_state( + self, claimed_generation: Optional[int] = None + ) -> _LoopState: """Gets or creates the state for the current event loop. + Args: + claimed_generation: The lifecycle generation the caller claimed + BEFORE its own awaits (setup passes the generation captured by + `_ensure_started`). Without it, a setup blocked ahead of this + call sampled the post-shutdown generation on resume and + published a writer that shutdown's snapshot could never see + (#6360 review round 7 P1-5). + Returns: The loop-specific state object containing clients and processors. """ @@ -3057,78 +4160,166 @@ async def _get_loop_state(self) -> _LoopState: 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. + # it (#6360 review round 5 P1-2). + raise RuntimeError("BigQuery plugin is shutting down.") + # Captured before any await: a shutdown() that starts (and even + # completes) while the writer below is being built bumps the + # generation, and the publication guard rechecks it — otherwise the + # new processor lands in the dict AFTER shutdown's snapshot/clear and + # leaks past close() (#6360 review round 6 P1-2). + generation = ( + claimed_generation + if claimed_generation is not None + else self._generation + ) + if self._generation != generation: raise RuntimeError("BigQuery plugin is shutting down.") self._cleanup_stale_loop_states() - # .get() rather than a membership test followed by indexing: a concurrent - # shutdown() clearing the dict under the guard between the two steps would - # otherwise raise KeyError here. - existing = self._loop_state_by_loop.get(loop) - if existing is not None: - return existing - - # grpc.aio clients are loop-bound, so we create one per event loop. - - def get_credentials() -> google.auth.credentials.Credentials: - creds, _ = google.auth.default(scopes=[_CLOUD_PLATFORM_SCOPE]) - return creds - - if self._credentials is None: - self._credentials = await loop.run_in_executor( - self._executor, get_credentials + detached_state: Optional[_LoopState] = None + detached_rows = 0 + detached_reason = "shutdown_timeout" + with self._loop_states_guard: + state = self._loop_state_by_loop.get(loop) + if state is not None: + processor = state.batch_processor + # Production entries always contain a real BatchProcessor. Keeping + # non-production stand-ins opaque also avoids treating truthy mock + # attributes as lifecycle flags in compatibility tests. + if not isinstance(processor, BatchProcessor): + return state + worker = processor._batch_processor_task + if worker is not None and not worker.done(): + if processor._shutdown: + # It is terminal for admission but still owns a live worker. + # Returning it loses rows; detaching/closing it races its drain. + raise _LoopStateAdmissionAbortedError( + "BigQuery writer is still shutting down." + ) + return state + + # A missing/done worker can never consume another appended row. + # Claim + fold under the same canonical guard order used by shutdown + # so concurrent stats readers see the state either live or folded, + # never both/neither. Identity ownership makes this single-winner. + detached_state = self._loop_state_by_loop.pop(loop) + # Prevent the old atexit registration from trying to run a second, + # blocking close over a processor whose rows are accounted below. + processor._shutdown = True + detached_reason = ( + "shutdown_cancelled" + if worker is not None and worker.cancelled() + else "shutdown_timeout" + ) + queue = processor._queue + sentinels = processor._sentinel_count + detached_rows = max(0, queue.qsize() - sentinels) + with self._drop_counts_guard: + for reason, count in processor.get_drop_stats().items(): + self._local_drop_counts[reason] = ( + self._local_drop_counts.get(reason, 0) + count + ) + if detached_rows: + self._local_drop_counts[detached_reason] = ( + self._local_drop_counts.get(detached_reason, 0) + detached_rows + ) + + if detached_rows: + logger.warning( + "%d queued row(s) belonged to a terminal BigQuery writer (%s).", + detached_rows, + detached_reason, + ) + # Structured ownership: the claimant that removed a terminal state also + # owns its bounded transport close. A detached fire-and-forget task left a + # warning/leak window whenever fresh construction raised before the task + # was retrieved. The finally runs on success, failure, and cancellation; + # on success the replacement is published before this await so concurrent + # callers share it instead of building another writer. + try: + # grpc.aio clients are loop-bound, so we create one per event loop. + def get_credentials() -> google.auth.credentials.Credentials: + creds, _ = google.auth.default(scopes=[_CLOUD_PLATFORM_SCOPE]) + return creds + + if self._credentials is None: + self._credentials = await loop.run_in_executor( + self._executor, get_credentials + ) + quota_project_id = getattr(self._credentials, "quota_project_id", None) + options = ( + client_options.ClientOptions(quota_project_id=quota_project_id) + if quota_project_id + else None ) - quota_project_id = getattr(self._credentials, "quota_project_id", None) - options = ( - client_options.ClientOptions(quota_project_id=quota_project_id) - if quota_project_id - else None - ) - user_agents = [f"google-adk-bq-logger/{__version__}"] - if self._visual_builder: - user_agents.append(f"google-adk-visual-builder/{__version__}") + user_agents = [f"google-adk-bq-logger/{__version__}"] + if self._visual_builder: + user_agents.append(f"google-adk-visual-builder/{__version__}") - client_info = gapic_client_info.ClientInfo(user_agent=" ".join(user_agents)) + client_info = gapic_client_info.ClientInfo( + user_agent=" ".join(user_agents) + ) - write_client = BigQueryWriteAsyncClient( - credentials=self._credentials, - client_info=client_info, - client_options=options, - ) + write_client = BigQueryWriteAsyncClient( + credentials=self._credentials, + client_info=client_info, + client_options=options, + ) - if not self._write_stream_name: - self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" - - batch_processor = BatchProcessor( - write_client=write_client, - arrow_schema=self.arrow_schema, - write_stream=self._write_stream_name, - batch_size=self.config.batch_size, - flush_interval=self.config.batch_flush_interval, - retry_config=self.config.retry_config, - queue_max_size=self.config.queue_max_size, - shutdown_timeout=self.config.shutdown_timeout, - ) - await batch_processor.start() + if not self._write_stream_name: + self._write_stream_name = f"projects/{self.project_id}/datasets/{self.dataset_id}/tables/{self.table_id}/_default" - state = _LoopState(write_client, batch_processor) - with self._loop_states_guard: - # Re-check under the guard: shutdown() may have started after the early - # _is_shutting_down check above. Publishing now would leak this live - # writer/processor past shutdown, so back out and tear it down instead. - published = not self._is_shutting_down - if published: - self._loop_state_by_loop[loop] = state - if not published: try: - await batch_processor.shutdown(timeout=self.config.shutdown_timeout) - except Exception: - pass - raise RuntimeError("BigQuery plugin is shutting down.") + batch_processor = BatchProcessor( + write_client=write_client, + arrow_schema=self.arrow_schema, + write_stream=self._write_stream_name, + batch_size=self.config.batch_size, + flush_interval=self.config.batch_flush_interval, + retry_config=self.config.retry_config, + queue_max_size=self.config.queue_max_size, + shutdown_timeout=self.config.shutdown_timeout, + ) + except BaseException: + # The write client already exists but no _LoopState can own it yet. + await self._close_write_transport(write_client) + raise + state = _LoopState(write_client, batch_processor) + try: + await batch_processor.start() + except BaseException: + # start() may create then fail/cancel a worker. Keep the fresh client + # under structured ownership as well; the bounded helper retrieves + # either sync or async transport-close outcomes. + await self._close_detached_loop_transport(state) + raise - atexit.register(self._atexit_cleanup, weakref.proxy(batch_processor)) + with self._loop_states_guard: + invalidated = self._is_shutting_down or self._generation != generation + if not invalidated: + self._loop_state_by_loop[loop] = state + if invalidated: + # shutdown() ran during construction; its snapshot cannot include + # this writer, so publishing it would leave a live processor and + # open transport behind after close() returns (#6360 review round 6 + # P1-2). Tear the fresh instances down instead of publishing. + try: + try: + await batch_processor.shutdown(timeout=self.config.shutdown_timeout) + except Exception: + logger.warning( + "Could not shut down writer created during shutdown.", + exc_info=True, + ) + finally: + await self._close_detached_loop_transport(state) + raise RuntimeError("BigQuery plugin is shutting down.") - return state + atexit.register(self._atexit_cleanup, weakref.proxy(batch_processor)) + return state + finally: + if detached_state is not None: + await self._close_detached_loop_transport(detached_state) async def flush(self) -> None: """Flushes any pending events to BigQuery. @@ -3153,39 +4344,91 @@ def get_drop_stats(self) -> dict[str, int]: 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. + ``formatter_failed`` and ``content_parse_failed`` mean the row WAS + written with its content replaced by a sentinel; ``setup_unavailable``, + ``shutdown_race``, ``shutdown_timeout``, ``shutdown_cancelled``, and + ``stale_loop`` mean the row was never written. Counters persist + across shutdown and loop cleanup. Returns: Per-reason counts: plugin-level incidents plus every live loop processor's counters (dead processors are folded in at shutdown/cleanup time). """ - totals: dict[str, int] = dict(self._local_drop_counts) - for state in list(self._loop_state_by_loop.values()): - for reason, count in state.batch_processor.get_drop_stats().items(): - totals[reason] = totals.get(reason, 0) + count + # Both guards, in the canonical order (loop states, then counters): + # reading them separately let a state that was folded-but-not-yet- + # removed be added twice, and a popped-but-not-yet-folded state be + # missed (#6360 review round 8 P2-7). + with self._loop_states_guard, self._drop_counts_guard: + totals: dict[str, int] = dict(self._local_drop_counts) + for state in self._loop_state_by_loop.values(): + for reason, count in state.batch_processor.get_drop_stats().items(): + totals[reason] = totals.get(reason, 0) + count return totals - async def _lazy_setup(self, **kwargs: Any) -> None: - """Performs lazy initialization of BigQuery clients and resources.""" + async def _lazy_setup( + self, claimed_generation: Optional[int] = None, **kwargs: Any + ) -> None: + """Performs lazy initialization of BigQuery clients and resources. + + Args: + claimed_generation: The lifecycle generation claimed by the + owning `_ensure_started` before any await; forwarded to + `_get_loop_state` so a shutdown that completes mid-setup is + detected even when it finishes before the loop-state phase + begins (#6360 review round 7 P1-5). + """ if self._started: return loop = asyncio.get_running_loop() - if not self.client: - if self._executor is None: - self._executor = ThreadPoolExecutor(max_workers=1) + # The executor is needed beyond client construction (schema RPCs, GCS + # offloader): creating it only inside the client branch left + # _executor as None for the offloader when a client was already set + # (post-rebase mypy: GCSOffloader argument 3 expects a non-optional + # ThreadPoolExecutor — a latent runtime gap, not just typing). + executor = self._executor + if executor is None: + executor = ThreadPoolExecutor(max_workers=1) + self._executor = executor - self.client = await loop.run_in_executor( - self._executor, + if not self.client: + client_future: "ConcurrentFuture[Any]" = executor.submit( lambda: bigquery.Client( project=self.project_id, credentials=self._credentials, ), ) + try: + self.client = await asyncio.wrap_future(client_future) + except asyncio.CancelledError: + # Cancelling the await does not stop the constructor thread; the + # eventual client was silently discarded and its connection pool + # never closed (#6360 review round 8 P2-14). The close itself is + # dispatched to a fresh thread: when the future is ALREADY done, + # add_done_callback runs the callback synchronously in THIS + # (event-loop) thread, and a slow client.close() would extend the + # host's cancellation window (#6360 review round 9 P1-7). + def _close_eventual(f: "ConcurrentFuture[Any]") -> None: + try: + eventual = f.result() + except Exception: + return + + def _close() -> None: + try: + eventual.close() + except Exception: + pass + + threading.Thread( + target=_close, + name="bqaa-orphan-client-close", + daemon=True, + ).start() + + client_future.add_done_callback(_close_eventual) + raise self.full_table_id = f"{self.project_id}.{self.dataset_id}.{self.table_id}" if not self._schema: @@ -3195,9 +4438,9 @@ async def _lazy_setup(self, **kwargs: Any) -> None: # 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. Once _started is True, + # 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) + await loop.run_in_executor(executor, self._ensure_schema_exists) if not self.parser: self.arrow_schema = to_arrow_schema(self._schema) @@ -3223,7 +4466,7 @@ async def _lazy_setup(self, **kwargs: Any) -> None: self.offloader = GCSOffloader( self.project_id, self.config.gcs_bucket_name, - self._executor, + executor, storage_client=storage.Client( project=self.project_id, credentials=self._credentials ), @@ -3237,7 +4480,7 @@ async def _lazy_setup(self, **kwargs: Any) -> None: connection_id=self.config.connection_id, ) - await self._get_loop_state() + await self._get_loop_state(claimed_generation=claimed_generation) @staticmethod def _atexit_cleanup(batch_processor: "BatchProcessor") -> None: @@ -3279,6 +4522,7 @@ def _ensure_schema_exists(self) -> None: exists, missing columns are added automatically (additive only). A ``adk_schema_version`` label is written for governance. """ + assert self.client is not None # _lazy_setup creates it before calling. try: existing_table = self.client.get_table(self.full_table_id) if self.config.auto_schema_upgrade: @@ -3297,10 +4541,16 @@ def _ensure_schema_exists(self) -> None: try: self.client.create_table(tbl) except cloud_exceptions.Conflict: - # Another process created it concurrently — still usable. - pass + # Another process created it concurrently — but there is no + # guarantee it used a compatible schema. Re-fetch and run the same + # readiness path as a pre-existing table; any failure here + # propagates so _ensure_started keeps _started=False and retries + # (#6360 review round 6 P1-4). + existing_table = self.client.get_table(self.full_table_id) + if self.config.auto_schema_upgrade: + self._maybe_upgrade_schema(existing_table) except Exception as e: - # Fail setup: returning normally here used to let the + # 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. @@ -3314,7 +4564,7 @@ def _ensure_schema_exists(self) -> None: if self.config.create_views: self._create_analytics_views() except Exception as e: - # Fail setup: swallowing control-plane errors here let + # Fail setup (#6356 P1-4): swallowing control-plane errors here let # the plugin mark itself started against a missing/unready table. logger.error( "Error ensuring table %s is ready: %s", @@ -3328,6 +4578,7 @@ def _ensure_schema_exists(self) -> None: def _schema_fields_match( existing: list[bq_schema.SchemaField], desired: list[bq_schema.SchemaField], + path: tuple[str, ...] = (), ) -> tuple[ list[bq_schema.SchemaField], list[bq_schema.SchemaField], @@ -3351,16 +4602,27 @@ def _schema_fields_match( existing_field = existing_by_name.get(desired_field.name) if existing_field is None: new_fields.append(desired_field) - elif ( - desired_field.field_type == "RECORD" - and existing_field.field_type == "RECORD" - and desired_field.fields - ): + continue + + field_path = ".".join((*path, desired_field.name)) + existing_type = existing_field.field_type.upper() + desired_type = desired_field.field_type.upper() + existing_mode = existing_field.mode.upper() + desired_mode = desired_field.mode.upper() + if existing_type != desired_type or existing_mode != desired_mode: + raise ValueError( + "Incompatible BigQuery schema field " + f"{field_path!r}: existing={existing_type}/{existing_mode}, " + f"desired={desired_type}/{desired_mode}." + ) + + if desired_type == "RECORD" and desired_field.fields: # Recurse into nested RECORD fields. sub_new, sub_updated = ( BigQueryAgentAnalyticsPlugin._schema_fields_match( list(existing_field.fields), list(desired_field.fields), + (*path, desired_field.name), ) ) if sub_new or sub_updated: @@ -3459,11 +4721,13 @@ def _maybe_upgrade_schema(self, existing_table: bigquery.Table) -> None: 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. + # 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. + # 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. @@ -3539,73 +4803,248 @@ async def create_analytics_views(self) -> None: loop = asyncio.get_running_loop() await loop.run_in_executor(self._executor, self._create_analytics_views) + @staticmethod + def _schedule_remote_drain( + processor: "BatchProcessor", + target_loop: asyncio.AbstractEventLoop, + drain_timeout: float, + ) -> "ConcurrentFuture[Any]": + """Schedules ``processor.shutdown()`` on another event loop. + + The coroutine is created INSIDE the remote-loop callback: + run_coroutine_threadsafe creates it eagerly in the caller thread, and + caller-side cleanup then had to guess ownership — + ``ConcurrentFuture.cancel()`` can return True even after the remote + task started, so a caller-side ``coro.close()`` either finalized the + coroutine on the wrong thread or raised "coroutine already executing" + (#6360 review round 10 P1-4). Here nothing exists to leak until the + callback runs, and the callback hands ownership to a remote Task + atomically via ``set_running_or_notify_cancel()``. + """ + cf: "ConcurrentFuture[Any]" = ConcurrentFuture() + + def _callback() -> None: + if not cf.set_running_or_notify_cancel(): + return # cancelled before the callback ran; nothing was created + + coro = processor.shutdown(timeout=drain_timeout) + try: + task = target_loop.create_task(coro) + except Exception as exc: + # e.g. a custom task factory rejecting creation: the coroutine + # exists but was never scheduled — close it here or it leaks as + # never-awaited (#6360 review round 11 P2-4). + coro.close() + cf.set_exception(exc) + return + + def _transfer(t: "asyncio.Task[Any]") -> None: + if t.cancelled(): + cf.set_exception(asyncio.CancelledError()) + elif t.exception() is not None: + cf.set_exception(t.exception()) + else: + cf.set_result(t.result()) + + task.add_done_callback(_transfer) + + target_loop.call_soon_threadsafe(_callback) + return cf + async def shutdown(self, timeout: float | None = None) -> None: """Shuts down the plugin and releases resources. Args: timeout: Maximum time to wait for the queue to drain. """ - if self._is_shutting_down: + while True: + waiter: Optional["ConcurrentFuture[None]"] = None + with self._setup_guard: + # Atomic admission (#6360 review round 7 P1-6): checked OUTSIDE + # the lock, two threads could both observe False, both claim + # shutdown, tear down the same snapshot twice, and double-fold + # identical drop counters. Exactly one caller per generation gets + # past this point. + if self._is_shutting_down: + waiter = self._shutdown_future + else: + 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 + self._shutdown_future = ConcurrentFuture() + if waiter is None: + break # this caller owns the teardown below + # Coalesce on the active owner: returning early made a concurrent + # `await plugin.close()` claim completion microseconds into another + # caller's teardown (#6360 review round 9 P1-5). shield: this + # waiter's own cancellation must not cancel the shared future. + try: + await asyncio.shield(asyncio.wrap_future(waiter)) + except _ShutdownIncompleteError: + # The owner was cancelled or failed mid-teardown; returning now + # would claim success while state is still live (#6360 review + # round 10 P1-3). Retry ownership — as a LOOP, not recursion, so + # depth does not grow with the number of coalesced callers + # (#6360 review round 12 P1-3). + continue + except Exception: + pass return - with self._setup_guard: - self._is_shutting_down = True - # Invalidate any in-flight setup: its completion must not resurrect - # _started after this method returns. - self._generation += 1 - self._started = False t = timeout if timeout is not None else self.config.shutdown_timeout loop = asyncio.get_running_loop() - # Re-affirm the shutdown flag and snapshot the live states in one critical - # section under _loop_states_guard -- the same guard _get_loop_state() - # holds for its publication re-check of _is_shutting_down. This serializes - # shutdown against a concurrent publisher without relying on GIL atomicity - # (correct under free-threaded builds too): the publisher either runs - # first, so its state is in this snapshot and gets drained, or observes - # the flag afterward and backs out. Snapshotting also avoids the - # "dictionary changed size during iteration" error from iterating the live - # dict. + # 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: - self._is_shutting_down = True states_snapshot = dict(self._loop_state_by_loop) + teardown_completed = False + teardown_error: Optional[BaseException] = None + retained_remote_drains = 0 try: # Correct Multi-Loop Shutdown: # 1. Shutdown current loop's processor directly. + drained: list[asyncio.AbstractEventLoop] = [] if loop in states_snapshot: await states_snapshot[loop].batch_processor.shutdown(timeout=t) - - # 1b. Drain batch processors on other (non-current) loops. + drained.append(loop) + + # 1b. Drain batch processors on other (non-current) loops. The + # wrapped futures are AWAITED, not .result()-ed: the synchronous + # wait blocked this event loop, so a host asyncio.timeout() around + # close() could never fire and the delay was paid serially per + # remote loop (#6360 review round 8 P1-6). One shared deadline + # covers all remote drains; unfinished ones are cancelled and their + # states left in place for a retry. + remote: list[ + tuple[ + asyncio.AbstractEventLoop, + "ConcurrentFuture[Any]", + "asyncio.Future[Any]", + ] + ] = [] for other_loop, state in states_snapshot.items(): if other_loop is loop: continue if other_loop.is_closed(): - # A closed loop cannot be driven to drain; count its unwritten - # queued rows so the loss is recorded before clear() below, - # mirroring _cleanup_stale_loop_states(). - stale_rows = self._count_unwritten_queued_rows(state) + # No drain is possible on a closed loop, and its queued rows + # are NOT guaranteed to have been counted — the state can enter + # this snapshot before _cleanup_stale_loop_states() ever ran + # (#6360 review round 9 P1-6). Claim, fold, and count the + # queue loss in ONE single-winner transition: counting from + # the snapshot without ownership double-counted rows that a + # concurrent stale cleanup had already claimed (#6360 review + # round 10 P2-5). + stale_rows = 0 + with self._loop_states_guard, self._drop_counts_guard: + owned = self._loop_state_by_loop.get(other_loop) is state + if owned: + del self._loop_state_by_loop[other_loop] + for ( + reason, + count, + ) in state.batch_processor.get_drop_stats().items(): + self._local_drop_counts[reason] = ( + self._local_drop_counts.get(reason, 0) + count + ) + queue = getattr(state.batch_processor, "_queue", None) + sentinels = getattr(state.batch_processor, "_sentinel_count", 0) + if isinstance(queue, asyncio.Queue): + if not isinstance(sentinels, int): + sentinels = 0 + stale_rows = max(0, queue.qsize() - sentinels) + if stale_rows: + self._local_drop_counts["stale_loop"] = ( + self._local_drop_counts.get("stale_loop", 0) + stale_rows + ) 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 during shutdown.", + "%d queued row(s) lost with closed loop %s.", stale_rows, id(other_loop), ) continue try: - future = asyncio.run_coroutine_threadsafe( - state.batch_processor.shutdown(timeout=t), - other_loop, - ) - future.result(timeout=t) + cf = self._schedule_remote_drain(state.batch_processor, other_loop, t) except Exception: + # e.g. the loop closed between the is_closed() check and + # call_soon_threadsafe(). The state stays live, so teardown is + # NOT complete — without counting it, both the owner and + # coalesced waiters reported success over live state (#6360 + # review round 14 P1-3). + retained_remote_drains += 1 logger.warning( "Could not drain batch processor on loop %s", other_loop, ) - - # 2. Close clients for all states - for state in states_snapshot.values(): + continue + remote.append((other_loop, cf, asyncio.wrap_future(cf))) + if remote: + try: + done_set, pending = await asyncio.wait( + [wrapper for _, _, wrapper in remote], timeout=t + ) + except asyncio.CancelledError: + # Host cancellation mid-wait: release every remote handle. The + # remote callback creates the coroutine itself, so a + # successfully cancelled concurrent future means nothing was + # (or ever will be) created (#6360 review round 10 P1-4). + for _, cf, wrapper in remote: + cf.cancel() + wrapper.cancel() + raise + del done_set + for other_loop, cf, wrapper in remote: + if wrapper in pending: + retained_remote_drains += 1 + # If the remote callback has not run yet this prevents the + # task from ever being created; if it HAS run, the running + # drain simply continues remotely, bounded by its own + # timeout, and the state is retained (#6360 review round 10 + # P1-4). + cf.cancel() + wrapper.cancel() + logger.warning( + "Batch processor drain on loop %s did not finish within" + " %.1fs; its state is retained for a retried close.", + other_loop, + t, + ) + continue + if wrapper.cancelled(): + retained_remote_drains += 1 + logger.warning( + "Batch processor drain on loop %s was cancelled; its" + " state is retained for a retried close.", + other_loop, + ) + continue + # Retrieve the result: an unchecked failed drain both leaked + # "exception was never retrieved" and claimed/folded the state + # as if it had succeeded, silently abandoning its queued rows + # (#6360 review round 9 P1-4). Only clean completions claim. + exc = wrapper.exception() + if exc is not None: + retained_remote_drains += 1 + logger.warning( + "Batch processor drain on loop %s failed (%s); its state" + " is retained for a retried close.", + other_loop, + type(exc).__name__, + ) + continue + drained.append(other_loop) + + # 2/3. For every DRAINED state: close its transport, then claim it + # out of the live dict and fold its counters in one atomic + # transition — fold-then-clear let get_drop_stats() add the same + # still-live processor again, and stale-loop cleanup could fold a + # snapshotted state a second time (#6360 review round 8 P2-7). + # States whose drain did not finish stay live (retry ownership). + for state_loop in drained: + state = states_snapshot[state_loop] if state.write_client and getattr( state.write_client, "transport", None ): @@ -3613,42 +5052,124 @@ async def shutdown(self, timeout: float | None = None) -> None: await state.write_client.transport.close() except Exception: pass - - # Fold processor drop counters into the persistent plugin-level - # counters before discarding loop state, so get_drop_stats() keeps - # reporting losses after shutdown. - 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() + with self._loop_states_guard, self._drop_counts_guard: + if self._loop_state_by_loop.get(state_loop) is state: + del self._loop_state_by_loop[state_loop] + for reason, count in state.batch_processor.get_drop_stats().items(): + self._local_drop_counts[reason] = ( + self._local_drop_counts.get(reason, 0) + count + ) + # else: stale-loop cleanup already claimed and folded it. # The parser/offloader hold the (now terminated) executor; keeping # them makes the first post-restart GCS upload raise "cannot - # schedule new futures after shutdown". - self.offloader = None + # schedule new futures after shutdown" (#6360 review round 5 P2). + # The offloader's plugin-owned storage.Client is CLOSED, not just + # dropped (#6360 review round 10 P2-6) — off-loop, under budget. + offloader, self.offloader = self.offloader, None self.parser = None + storage_client = getattr(offloader, "client", None) if offloader else None + if storage_client is not None: + try: + await asyncio.wait_for( + loop.run_in_executor(None, storage_client.close), timeout=t + ) + except Exception: + pass - if self.client: - if self._executor: - executor = self._executor - await loop.run_in_executor(None, lambda: executor.shutdown(wait=True)) - self._executor = None - self.client = None + # The executor is shut down INDEPENDENTLY of the client: a + # cancelled setup could leave a live executor with client=None, and + # the nested check leaked it past close() (#6360 review round 8 + # P2-14). Non-blocking: waiting would stall close() behind a + # slow/blocked constructor job in the pool; pending jobs are + # cancelled, and a still-running constructor finishes in the + # background (its orphaned client is closed by _lazy_setup's + # done-callback). + if self._executor: + self._executor.shutdown(wait=False, cancel_futures=True) + self._executor = None + # Close (not just drop) the shared BigQuery client: discarding the + # reference leaked its HTTP transport, while the aborted-setup path + # already closed the same resource (#6360 review round 9 P2-11). + # Off-loop and bounded by the shutdown budget. + client, self.client = self.client, None + if client is not None: + try: + await asyncio.wait_for( + loop.run_in_executor(None, client.close), timeout=t + ) + except Exception: + pass + if retained_remote_drains: + # An incompletely drained remote state means live processors and + # possibly queued rows survive this close: reporting success let + # both the owner and coalesced waiters return normally over live + # state (#6360 review round 13 P1-3, extending the round-12 + # honest-completion contract to remote drains). + raise _ShutdownIncompleteError( + f"{retained_remote_drains} remote drain(s) did not complete;" + " their states are retained for a retried close." + ) + teardown_completed = True except Exception as e: + # teardown_completed stays False: reporting success here let both + # the owner and coalesced waiters return normally while the loop + # state was still live (#6360 review round 11 P1-3). Waiters + # receive _ShutdownIncompleteError and retry ownership; the OWNER + # re-raises after the finally so its caller sees the failure too — + # PluginManager.close() aggregates plugin close failures (#6360 + # review round 12 P1-3). + teardown_error = e logger.error("Error during shutdown: %s", e, exc_info=True) - self._is_shutting_down = False - self._started = False + finally: + # Cancellation-safe reset: PluginManager's close timeout cancels this + # coroutine, and asyncio.CancelledError is a BaseException that the + # handler above does not (and must not) swallow. Without the finally, + # a cancelled shutdown left _is_shutting_down=True forever, so the + # re-entry guard turned every later close() into a no-op and retained + # state could never be cleaned up (#6360 review round 6 P1-1). Any + # loop states not yet drained stay in _loop_state_by_loop, so a + # retried shutdown() re-snapshots and finishes the job; the + # cancellation itself propagates to the caller unchanged. + # ONE guarded transition for the admission flag, lifecycle flags, + # and the completion-future swap: resetting _is_shutting_down + # before taking the guard let a new caller claim shutdown and + # install ITS future in the gap, after which this owner resolved + # the wrong future and a third caller could observe + # _is_shutting_down=True with no future and fall into overlapping + # teardown (#6360 review round 11 P1-2). + with self._setup_guard: + self._is_shutting_down = False + self._started = False + completion, self._shutdown_future = self._shutdown_future, None + # Wake coalesced callers. Success is only reported when teardown + # actually ran to completion: resolving unconditionally let an + # uncancelled waiter return from close() while the owner was + # cancelled mid-teardown and state was still live (#6360 review + # round 10 P1-3). On the incomplete path waiters retry ownership. + if completion is not None and not completion.done(): + if teardown_completed: + completion.set_result(None) + else: + completion.set_exception( + _ShutdownIncompleteError( + "Owning shutdown did not complete teardown." + ) + ) + if teardown_error is not None: + # The owning caller must not report success over live state (#6360 + # review round 12 P1-3). + raise teardown_error def __getstate__(self) -> dict[str, Any]: """Custom pickling to exclude non-picklable runtime objects.""" state = self.__dict__.copy() state["_setup_guard"] = None state["_setup_future"] = None + state["_shutdown_future"] = None state["_generation"] = 0 state["_loop_states_guard"] = None + state["_drop_counts_guard"] = None state["client"] = None state["_loop_state_by_loop"] = {} state["_write_stream_name"] = None @@ -3677,11 +5198,13 @@ def __setstate__(self, state: dict[str, Any]) -> None: self.__dict__.update(state) self._setup_guard = threading.Lock() self._setup_future = None + self._shutdown_future = None self._generation = 0 self._loop_states_guard = threading.Lock() + self._drop_counts_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. + # otherwise skip the write loop silently (#6360 review round 2 P2-7). _validate_runtime_config(self.config) def _reset_runtime_state(self) -> None: @@ -3727,8 +5250,10 @@ def _reset_runtime_state(self) -> None: # Clear all runtime state. self._setup_guard = threading.Lock() self._setup_future = None + self._shutdown_future = None self._generation = 0 self._loop_states_guard = threading.Lock() + self._drop_counts_guard = threading.Lock() self.client = None self._loop_state_by_loop = {} self._write_stream_name = None @@ -3743,16 +5268,19 @@ def _reset_runtime_state(self) -> None: self._init_pid = os.getpid() def _count_local_drop(self, reason: str) -> None: - """Counts a row lost before/outside any BatchProcessor.""" - self._local_drop_counts[reason] = self._local_drop_counts.get(reason, 0) + 1 + """Counts a row lost before/outside any BatchProcessor (#6356 P2).""" + with self._drop_counts_guard: + 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. PluginManager's outer close - timeout (5s) may cancel this + 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. """ @@ -3770,20 +5298,30 @@ async def __aexit__( ) -> None: await self.shutdown() - async def _ensure_started(self, **kwargs: Any) -> None: + async def _ensure_started(self, **kwargs: Any) -> str: """Ensures that the plugin is started and initialized. - Setup failures no longer poison the plugin permanently: + 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. + + Returns: + A structured outcome — ``"ok"``, ``"disabled"``, ``"failed"``, or + ``"aborted"`` (shutdown crossed the attempt). This method never + counts a lost row itself: it is also called from non-row paths + (Runner start, ``__aenter__``), which produced phantom + ``shutdown_race`` counts, while the real row owner counted the + same incident a second time as ``setup_unavailable`` (#6360 + review round 8 P2-9). The row owner counts exactly one loss + based on this outcome. """ - # Disabled mode must have zero side effects: no ADC lookup, + # 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 + return "disabled" # _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 @@ -3794,9 +5332,10 @@ async def _ensure_started(self, **kwargs: Any) -> None: if self._init_pid != 0 and os.getpid() != self._init_pid: self._reset_runtime_state() if self._started: - return + return "ok" - # Cross-loop coalescing of the SHARED initialization: _lazy_setup mutates process-wide state (client, + # 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 @@ -3809,7 +5348,7 @@ async def _ensure_started(self, **kwargs: Any) -> None: is_owner = False with self._setup_guard: if self._started: - return + return "ok" if self._setup_future is not None: setup_future = self._setup_future elif ( @@ -3817,7 +5356,7 @@ async def _ensure_started(self, **kwargs: Any) -> None: and time.monotonic() < self._setup_retry_at ): # Still inside the backoff window from a previous failure. - return + return "failed" else: setup_future = ConcurrentFuture() self._setup_future = setup_future @@ -3830,37 +5369,89 @@ async def _ensure_started(self, **kwargs: Any) -> None: 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. The waiter itself still observes its own + # 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 _SetupAbortedError: + return "aborted" 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 + # degrade the same way the owner does (loss counted by the row + # owner from this outcome). + return "failed" + return "ok" if self._started else "failed" try: - await self._lazy_setup(**kwargs) + await self._lazy_setup(claimed_generation=claimed_generation, **kwargs) except asyncio.CancelledError: # Owner cancelled mid-setup: without this, the pending future was - # never finalized and every later _ensure_started waited forever. - # Clear the rendezvous, wake waiters - # with an ordinary aborted error, then re-raise the cancellation. + # never finalized and every later _ensure_started waited forever + # (#6360 review round 5 P1-1). Release partial resources that need + # no await (the executor keeps running its current job; the + # eventual client is closed by _lazy_setup's done-callback, #6360 + # review round 8 P2-14), clear the rendezvous, wake waiters with an + # ordinary aborted error, then re-raise the cancellation. + executor, self._executor = self._executor, None + if executor is not None and self.client is None: + executor.shutdown(wait=False) + elif executor is not None: + self._executor = executor # a live client still uses it + self.offloader = None + self.parser = None with self._setup_guard: self._setup_future = None if not setup_future.done(): setup_future.set_exception( - RuntimeError("BigQuery plugin setup aborted: owner cancelled.") + _SetupAbortedError( + "BigQuery plugin setup aborted: owner cancelled." + ) ) raise - except Exception as e: + except _LoopStateAdmissionAbortedError as e: + # A retained processor with _shutdown=True and a live worker is a + # lifecycle admission race, not a service/setup failure. Keep shared + # clients intact, avoid poisoning exponential backoff, and wake every + # coalesced caller with the same structured aborted outcome. The row + # owner (and only the row owner) converts that outcome to shutdown_race. 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 + if not setup_future.done(): + setup_future.set_exception(e) + return "aborted" + except Exception as e: + aborted = False + with self._setup_guard: + if self._generation != claimed_generation: + # shutdown() completed while setup was blocked; the failure is + # the abort itself, not a service error, so it must not poison + # the backoff window — and the partially created resources must + # be released (#6360 review round 7 P1-5). + aborted = True + else: + 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 + if aborted: + # The rendezvous stays claimed until teardown completes: clearing + # it first let a new-generation setup finish while the old owner + # was paused, after which this teardown destroyed the NEW + # client/parser/loop state and the plugin wedged with + # _started=True and no resources (#6360 review round 8 P1-5). + try: + await self._teardown_aborted_setup() + finally: + with self._setup_guard: + self._setup_future = None + if not setup_future.done(): + setup_future.set_exception( + _SetupAbortedError( + "BigQuery plugin setup aborted: shutdown during setup." + ) + ) + return "aborted" logger.error( "Failed to initialize BigQuery Plugin (attempt %d, next" " retry in %.0fs): %s", @@ -3870,14 +5461,14 @@ async def _ensure_started(self, **kwargs: Any) -> None: ) if not setup_future.done(): setup_future.set_exception(e) + return "failed" 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. + # _started after shutdown returned (#6360 review round 5 P1-2). aborted = True - self._setup_future = None else: self._started = True self._startup_error = None @@ -3888,17 +5479,88 @@ async def _ensure_started(self, **kwargs: Any) -> None: # the rest of this instance's lifetime. if self._init_pid == 0: self._init_pid = os.getpid() - if not setup_future.done(): - if aborted: + if not aborted: + if not setup_future.done(): + setup_future.set_result(None) + return "ok" + # Setup fully succeeded but lost the generation race: everything it + # created outlives a shutdown that already returned — release it + # (#6360 review round 7 P1-5), holding the rendezvous until the + # teardown completes (#6360 review round 8 P1-5). + try: + await self._teardown_aborted_setup() + finally: + with self._setup_guard: + self._setup_future = None + if not setup_future.done(): setup_future.set_exception( - RuntimeError( + _SetupAbortedError( "BigQuery plugin setup aborted: shutdown during setup." ) ) - else: - setup_future.set_result(None) - if aborted: - self._count_local_drop("shutdown_race") + return "aborted" + + async def _teardown_aborted_setup(self) -> None: + """Releases every resource created by a setup that crossed a shutdown. + + A setup attempt that lost the generation race used to count the loss + and stop, leaving the freshly created shared client, executor, + parser/offloader, and any published loop state alive on a plugin + whose shutdown() had already returned (#6360 review round 7 P1-5). + Callers must hold the setup rendezvous (_setup_future) for the whole + teardown so no new-generation setup can publish resources this method + would then destroy (#6360 review round 8 P1-5); the started-check is + a second line of defense. + """ + if self._started: + # A newer-generation setup owns the current resources. + return + try: + loop = asyncio.get_running_loop() + except RuntimeError: + loop = None + state = None + if loop is not None: + with self._loop_states_guard: + state = self._loop_state_by_loop.pop(loop, None) + if state is not None: + try: + await state.batch_processor.shutdown( + timeout=self.config.shutdown_timeout + ) + except Exception: + pass + transport = getattr(state.write_client, "transport", None) + if transport: + try: + await transport.close() + except Exception: + pass + offloader, self.offloader = self.offloader, None + self.parser = None + storage_client = getattr(offloader, "client", None) if offloader else None + if storage_client is not None: + # Close the owned GCS client too (#6360 review round 10 P2-6). + try: + await asyncio.get_running_loop().run_in_executor( + None, storage_client.close + ) + except Exception: + pass + client, self.client = self.client, None + executor, self._executor = self._executor, None + if client is not None: + try: + await asyncio.get_running_loop().run_in_executor(None, client.close) + except Exception: + pass + if executor is not None: + try: + await asyncio.get_running_loop().run_in_executor( + None, lambda: executor.shutdown(wait=True) + ) + except Exception: + pass @staticmethod def _resolve_ids( @@ -4283,29 +5945,106 @@ async def _log_event( return if not self._started: - await self._ensure_started() + outcome = await self._ensure_started() + if outcome == "disabled": + return if not self._started: - # Setup unavailable (failed and inside its retry backoff): the row - # is lost — record it so the loss is observable. - self._count_local_drop("setup_unavailable") + # The row is lost — record exactly ONE loss, classified by the + # structured setup outcome: "aborted" (shutdown crossed the + # attempt) is a shutdown_race, anything else is setup + # unavailability. _ensure_started itself never counts, so + # non-row entry points no longer produce phantom counts and one + # raced event no longer counts twice (#6360 review round 8 P2-9). + self._count_local_drop( + "shutdown_race" if outcome == "aborted" else "setup_unavailable" + ) return if event_data is None: event_data = EventData() + # Error diagnostics bypass the ordinary attributes tree: error_message is + # a dedicated column and agent/run tracebacks live in raw content. Apply + # one bounded, fail-closed boundary here so every current and future error + # producer receives the same privacy contract before formatter/parser row + # assembly. Ordinary safe messages remain byte-for-byte unchanged. + if event_data.error_message is not None: + try: + safe_error, error_content_lost = _sanitize_sensitive_text( + event_data.error_message, self.config.max_content_length + ) + except Exception: + safe_error, error_content_lost = ( + "[REDACTED_SENSITIVE_TEXT]", + True, + ) + event_data.error_message = safe_error + is_truncated = is_truncated or error_content_lost + if event_type in ("AGENT_ERROR", "INVOCATION_ERROR") and isinstance( + raw_content, collections.abc.Mapping + ): + try: + error_traceback = raw_content.get("error_traceback") + if isinstance(error_traceback, str): + safe_traceback, traceback_content_lost = _sanitize_sensitive_text( + error_traceback, self.config.max_content_length + ) + raw_content = dict(raw_content) + raw_content["error_traceback"] = safe_traceback + is_truncated = is_truncated or traceback_content_lost + except Exception: + raw_content = {"error_traceback": "[REDACTED_SENSITIVE_TEXT]"} + is_truncated = True + timestamp = datetime.now(timezone.utc) if self.config.content_formatter: try: - raw_content = self.config.content_formatter(raw_content, event_type) - except Exception as e: - # Fail CLOSED: the formatter is a redaction/privacy + formatted = self.config.content_formatter(raw_content, event_type) + if isinstance(formatted, str): + if type(formatted) is not str: + # Normalize str subclasses to the exact built-in (#6360 + # round 9 P1-1 / round 7 P1-3). + formatted = str.__str__(formatted) + elif formatted is not None and not ( + # Every shape the parser handles NATIVELY: identity and + # conditional formatters legitimately return these, and the + # round-9 str/Content/None-only gate destroyed untransformed + # LlmRequest/dict/list events (#6360 review round 10 P1-2). + # Model shapes require the EXACT class: a subclass can + # override an attribute the parser reads OUTSIDE this + # boundary and raise a payload-bearing exception into the + # safe callback's traceback log (#6360 review round 12 + # P1-2). dict/list subclasses stay isinstance-based — the + # parser routes them through the hardened recursive + # sanitizer, whose protocol boundary already fails closed. + type(formatted) in (types.Content, types.Part, LlmRequest) + or isinstance(formatted, (dict, list)) + ): + # The formatter is typed Any: a non-native result would reach + # the parser's unconditional str(content) fallback OUTSIDE this + # fail-closed boundary, where a payload-controlled __str__ can + # republish the original content or raise into the safe + # callback's traceback log (#6360 review round 9 P1-1). The + # message is CONSTANT: even a class NAME can be payload-derived + # via type(name, ...) (#6360 review round 10 P2-7). + logger.warning( + "Content formatter returned an unsupported result type for" + " event %s; writing sentinel instead of original content.", + event_type, + ) + formatted = _FORMATTER_FAILED_SENTINEL + self._count_local_drop("formatter_failed") + raw_content = formatted + except Exception: + # Fail CLOSED (#6356 P1-1): the formatter is a redaction/privacy # boundary, so its failure must never fall back to the unformatted - # payload. Log only the exception CLASS — the message or a - # traceback (exc_info) could embed the protected content itself. + # payload. The log message is CONSTANT — the exception message and + # traceback can embed the protected content, and even the class + # NAME can be payload-derived via type(name, ...) (#6360 review + # round 10 P2-7). logger.warning( - "Content formatter (%s) failed for event %s; writing sentinel" + "Content formatter failed for event %s; writing sentinel" " instead of original content.", - type(e).__name__, event_type, ) raw_content = _FORMATTER_FAILED_SENTINEL @@ -4330,12 +6069,53 @@ async def _log_event( else: # 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. - content_json, content_parts, parser_truncated = await self.parser.parse( - raw_content, - trace_id=trace_id or "no_trace", - span_id=span_id or "no_span", - ) + # GCS object paths (#6356 P1-3). + try: + content_json, content_parts, parser_truncated = await self.parser.parse( + raw_content, + trace_id=trace_id or "no_trace", + span_id=span_id or "no_span", + ) + # Normalize the parser OUTPUT to strictly JSON-native values + # inside the same boundary (#6360 review round 14 P1-2): a nested + # hostile model can defer its failure PAST parse(), detonating in + # Arrow serialization's json.dumps/str fallback where + # _write_rows_with_retry logged the payload with a traceback and + # dropped the row as arrow_prep_failed. + content_json, norm_replaced_json = _normalize_json_native( + content_json, self.config.max_content_length + ) + # content_parts carry parser-BUILT metadata (GCS URIs, + # object_ref.details JSON) whose strings must stay intact; their + # payload text was already truncated by the parser itself, so + # only shape normalization applies (max_len=-1). + normalized_parts, norm_replaced_parts = _normalize_json_native( + content_parts, -1 + ) + content_parts = ( + normalized_parts if isinstance(normalized_parts, list) else [] + ) + parser_truncated = ( + parser_truncated or norm_replaced_json or norm_replaced_parts + ) + except Exception: + # Fail-closed, constant-log parse boundary (#6360 review round 13 + # P1-1): the top-level formatter gate cannot see NESTED hostile + # model subclasses (pydantic preserves them through normal + # construction), whose attribute accesses raise payload-bearing + # exceptions inside the parser. Escaping here reached + # _safe_callback's traceback log and dropped the whole row. + logger.warning( + "Content parsing failed for event %s; writing sentinel" + " instead of content.", + event_type, + ) + content_json, content_parts, parser_truncated = ( + "[CONTENT_PARSE_FAILED]", + [], + True, + ) + self._count_local_drop("content_parse_failed") is_truncated = is_truncated or parser_truncated latency_json = self._extract_latency(event_data) @@ -4349,7 +6129,7 @@ async def _log_event( meta_truncated = self._capture_custom_metadata(event_data, attributes) is_truncated = is_truncated or meta_truncated - # Final safety pass: sanitize the COMPLETE assembled + # 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 @@ -4521,6 +6301,15 @@ async def on_event_callback( """ callback_ctx = CallbackContext(invocation_context) + # A later before_model callback may short-circuit the model call. ADK + # intentionally skips every after_model callback in that case, so the + # llm_request span pushed by this plugin has no matching callback to pop + # it. The synthesized response is a non-partial event; close only that + # expected top span at this boundary. Streaming chunks stay attached to + # their live llm_request span until the final response callback. + if getattr(event, "partial", None) is not True: + TraceManager.pop_span(expected_kind="llm_request") + # --- State delta logging --- if event.actions.state_delta: await self._log_event( @@ -4906,6 +6695,12 @@ async def before_model_callback( {system_prompt}'. """ + # Defensive cleanup for a short-circuited request whose synthesized + # event was not observed (for example, an abnormal generator exit). + # expected_kind prevents this from disturbing the parent agent or + # invocation span. + TraceManager.pop_span(expected_kind="llm_request") + # 5. Attributes (Config & Tools) attributes: dict[str, Any] = {} tools_truncated = False @@ -5034,7 +6829,9 @@ async def after_model_callback( tfft = int((first_token - start_time) * 1000) # ACTUALLY pop the span - popped_span_id, duration = TraceManager.pop_span() + popped_span_id, duration = TraceManager.pop_span( + expected_kind="llm_request" + ) is_popped = True # If we popped, the span_id from get_current_span_and_parent() above is correct for THIS event @@ -5075,7 +6872,7 @@ async def on_model_error_callback( llm_request: The request that was sent to the model. error: The exception that occurred. """ - span_id, duration = TraceManager.pop_span() + span_id, duration = TraceManager.pop_span(expected_kind="llm_request") parent_span_id, _ = TraceManager.get_current_span_and_parent() await self._log_event( @@ -5258,10 +7055,6 @@ async def on_agent_error_callback( type(error), error, error.__traceback__ ) ) - max_len = self.config.max_content_length - if max_len > 0 and len(error_tb) > max_len: - error_tb = error_tb[:max_len] + "... [truncated]" - await self._log_event( "AGENT_ERROR", callback_context, @@ -5307,10 +7100,6 @@ async def on_run_error_callback( type(error), error, error.__traceback__ ) ) - max_len = self.config.max_content_length - if max_len > 0 and len(error_tb) > max_len: - error_tb = error_tb[:max_len] + "... [truncated]" - await self._log_event( "INVOCATION_ERROR", callback_ctx, diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 340e583723..82f7239795 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -14,14 +14,14 @@ from __future__ import annotations import asyncio -import collections import contextlib import dataclasses import json import logging import os +import sys import threading -from types import MappingProxyType +import time from unittest import mock from google.adk.agents import base_agent @@ -407,6 +407,10 @@ def test_recursive_smart_truncate_redaction(): "id_token": "eyJhb", "api_key": "AIza", "password": "my-password", + "private_key": "private-key-material", + "token": "generic-token", + "secret": "generic-secret", + "authorization": "Bearer credential", "safe_key": "safe-value", "temp:auth_state": "some-auth-state", "nested": { @@ -425,6 +429,10 @@ def test_recursive_smart_truncate_redaction(): assert truncated["id_token"] == "[REDACTED]" assert truncated["api_key"] == "[REDACTED]" assert truncated["password"] == "[REDACTED]" + assert truncated["private_key"] == "[REDACTED]" + assert truncated["token"] == "[REDACTED]" + assert truncated["secret"] == "[REDACTED]" + assert truncated["authorization"] == "[REDACTED]" assert truncated["safe_key"] == "safe-value" assert truncated["temp:auth_state"] == "[REDACTED]" assert truncated["nested"]["CLIENT_SECRET"] == "[REDACTED]" @@ -764,7 +772,7 @@ def error_formatter(content, event_type): log_entry = await _get_captured_event_dict_async( mock_write_client, dummy_arrow_schema ) - # Fail CLOSED: a raising formatter must never fall back + # 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"]) @@ -1817,7 +1825,9 @@ async def test_after_model_callback_text_response( prompt_token_count=10, total_token_count=15 ), ) - bigquery_agent_analytics_plugin.TraceManager.push_span(callback_context) + bigquery_agent_analytics_plugin.TraceManager.push_span( + callback_context, "llm_request" + ) await bq_plugin_inst.after_model_callback( callback_context=callback_context, llm_response=llm_response, @@ -3475,7 +3485,7 @@ async def test_parser_identity_not_mutated_per_call( ): """_log_event must NOT store request identity on the shared parser. - trace_id/span_id are passed per parse() call: mutating the + 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. """ @@ -5191,7 +5201,7 @@ def test_upgrade_adds_missing_columns(self): plugin = self._make_plugin(auto_schema_upgrade=True) existing = mock.MagicMock(spec=bigquery.Table) existing.schema = [ - bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"), ] existing.labels = {"other": "label"} plugin.client.get_table.return_value = existing @@ -5227,12 +5237,13 @@ 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. + 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 = [ - bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"), ] existing.labels = {} plugin.client.get_table.return_value = existing @@ -5240,6 +5251,35 @@ def test_upgrade_error_propagates_when_fields_missing(self): with pytest.raises(Exception, match="boom"): plugin._ensure_schema_exists() + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("existing_type", "existing_mode"), + [("STRING", "REQUIRED"), ("TIMESTAMP", "NULLABLE")], + ids=("type", "mode"), + ) + async def test_incompatible_existing_field_blocks_startup( + self, existing_type, existing_mode + ): + """Same-name fields with incompatible type/mode are not ready.""" + plugin = self._make_plugin(auto_schema_upgrade=True) + plugin.config.create_views = False + existing = mock.MagicMock(spec=bigquery.Table) + existing.schema = [ + bigquery.SchemaField("timestamp", existing_type, mode=existing_mode), + ] + existing.labels = {} + plugin.client.get_table.return_value = existing + + try: + outcome = await plugin._ensure_started() + assert outcome == "failed" + assert plugin._started is False + assert isinstance(plugin._startup_error, ValueError) + assert "timestamp" in str(plugin._startup_error) + plugin.client.update_table.assert_not_called() + finally: + await plugin.shutdown() + def test_upgrade_preserves_existing_columns(self): """Existing columns are never dropped or altered during upgrade.""" plugin = self._make_plugin(auto_schema_upgrade=True) @@ -5248,7 +5288,7 @@ def test_upgrade_preserves_existing_columns(self): custom_field = bigquery.SchemaField("my_custom_col", "STRING") existing = mock.MagicMock(spec=bigquery.Table) existing.schema = [ - bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"), bigquery.SchemaField("event_type", "STRING"), custom_field, ] @@ -5291,7 +5331,7 @@ def test_upgrade_from_older_version_label(self): plugin = self._make_plugin(auto_schema_upgrade=True) existing = mock.MagicMock(spec=bigquery.Table) existing.schema = [ - bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"), bigquery.SchemaField("event_type", "STRING"), ] # Simulate a table stamped with an older version. @@ -5322,7 +5362,7 @@ def test_upgrade_is_idempotent(self): # First call: table exists with old schema. existing = mock.MagicMock(spec=bigquery.Table) existing.schema = [ - bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"), ] existing.labels = {} plugin.client.get_table.return_value = existing @@ -5344,7 +5384,7 @@ def test_update_table_receives_schema_and_labels_fields(self): plugin = self._make_plugin(auto_schema_upgrade=True) existing = mock.MagicMock(spec=bigquery.Table) existing.schema = [ - bigquery.SchemaField("timestamp", "TIMESTAMP"), + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED"), ] existing.labels = {} plugin.client.get_table.return_value = existing @@ -5360,15 +5400,61 @@ def test_auto_schema_upgrade_defaults_to_true(self): config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig() assert config.auto_schema_upgrade is True - def test_create_table_conflict_is_ignored(self): - """Race condition (Conflict) during create_table is silently handled.""" + def test_create_table_conflict_refetches_concurrent_table(self): + """Conflict during create_table re-fetches the concurrently created + table instead of blindly trusting it (#6360 round 6 P1-4).""" plugin = self._make_plugin() - plugin.client.get_table.side_effect = cloud_exceptions.NotFound("not found") + existing = mock.MagicMock(spec=bigquery.Table) + existing.schema = plugin._schema + existing.labels = {} + plugin.client.get_table.side_effect = [ + cloud_exceptions.NotFound("not found"), + existing, + ] plugin.client.create_table.side_effect = cloud_exceptions.Conflict( "already exists" ) # Should not raise. plugin._ensure_schema_exists() + assert plugin.client.get_table.call_count == 2 + + def test_create_table_conflict_upgrades_incompatible_table(self): + """A concurrently created table missing required columns goes through + the normal upgrade path after Conflict (#6360 round 6 P1-4).""" + plugin = self._make_plugin(auto_schema_upgrade=True) + incompatible = mock.MagicMock(spec=bigquery.Table) + incompatible.schema = [ + bigquery.SchemaField("timestamp", "TIMESTAMP", mode="REQUIRED") + ] + incompatible.labels = {} + plugin.client.get_table.side_effect = [ + cloud_exceptions.NotFound("not found"), + incompatible, + ] + plugin.client.create_table.side_effect = cloud_exceptions.Conflict( + "already exists" + ) + plugin._ensure_schema_exists() + assert plugin.client.get_table.call_count == 2 + plugin.client.update_table.assert_called_once() + updated_names = { + f.name for f in plugin.client.update_table.call_args[0][0].schema + } + assert "event_type" in updated_names + + def test_create_table_conflict_refetch_failure_propagates(self): + """If the post-Conflict readiness check fails, setup must fail so + _ensure_started retries later (#6360 round 6 P1-4).""" + plugin = self._make_plugin(auto_schema_upgrade=True) + plugin.client.get_table.side_effect = [ + cloud_exceptions.NotFound("not found"), + cloud_exceptions.ServiceUnavailable("control plane down"), + ] + plugin.client.create_table.side_effect = cloud_exceptions.Conflict( + "already exists" + ) + with pytest.raises(cloud_exceptions.ServiceUnavailable): + plugin._ensure_schema_exists() class TestToolProvenance: @@ -6510,7 +6596,7 @@ async def test_create_analytics_views_ensures_started( await plugin.shutdown() def test_views_not_created_after_table_creation_failure(self): - """create_table failure raises (fail setup) and skips views.""" + """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") @@ -7741,6 +7827,31 @@ def test_nested_field_detected(self): assert "key" in sub_names assert "value" in sub_names + def test_nested_field_mode_mismatch_is_rejected(self): + """Nested same-name fields must match type and mode too.""" + plugin = self._make_plugin() + plugin._schema = [ + bigquery.SchemaField( + "metadata", + "RECORD", + fields=[bigquery.SchemaField("key", "STRING", mode="REQUIRED")], + ) + ] + existing = mock.MagicMock(spec=bigquery.Table) + existing.schema = [ + bigquery.SchemaField( + "metadata", + "RECORD", + fields=[bigquery.SchemaField("key", "STRING", mode="NULLABLE")], + ) + ] + existing.labels = {} + plugin.client.get_table.return_value = existing + + with pytest.raises(ValueError, match=r"metadata\.key"): + plugin._ensure_schema_exists() + plugin.client.update_table.assert_not_called() + def test_version_label_not_stamped_on_failure(self): """A failed update_table does not persist the version label.""" plugin = self._make_plugin() @@ -7758,7 +7869,7 @@ def test_version_label_not_stamped_on_failure(self): plugin.client.update_table.side_effect = Exception("network error") # Raises so setup is not marked ready against a table with missing - # fields. + # fields (#6360 review round 2 P1-2). with pytest.raises(Exception, match="network error"): plugin._ensure_schema_exists() @@ -7836,7 +7947,12 @@ async def test_other_loop_batch_processor_drained( mock_to_arrow_schema, mock_asyncio_to_thread, ): - """Shutdown drains batch_processor.shutdown on non-current loops.""" + """Shutdown drains batch_processor.shutdown on non-current loops. + + Uses a REAL second loop: since #6360 round 10 P1-4 the drain task is + created inside the remote loop's own callback (no + run_coroutine_threadsafe), so the drain must actually execute there. + """ plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( project_id=PROJECT_ID, dataset_id=DATASET_ID, @@ -7844,42 +7960,41 @@ async def test_other_loop_batch_processor_drained( ) await plugin._ensure_started() - # Create a mock "other" loop with a mock batch processor. - other_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop) - other_loop.is_closed.return_value = False - - mock_other_bp = mock.AsyncMock() - mock_other_write_client = mock.MagicMock() - mock_other_write_client.transport = mock.AsyncMock() + other_loop = asyncio.new_event_loop() + thread = threading.Thread(target=other_loop.run_forever, daemon=True) + thread.start() + try: + drain_thread_ids = [] - other_state = bigquery_agent_analytics_plugin._LoopState( - write_client=mock_other_write_client, - batch_processor=mock_other_bp, - ) - plugin._loop_state_by_loop[other_loop] = other_state + async def record_shutdown(timeout=None): + del timeout + drain_thread_ids.append(threading.get_ident()) - # Patch run_coroutine_threadsafe to verify it's called for - # the other loop's batch_processor. Close the coroutine arg - # to avoid "coroutine was never awaited" RuntimeWarning. - mock_future = mock.MagicMock() - mock_future.result.return_value = None + mock_other_bp = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + mock_other_bp.shutdown = record_shutdown + mock_other_bp.get_drop_stats = mock.MagicMock(return_value={}) + mock_other_write_client = mock.MagicMock() + mock_other_write_client.transport = mock.AsyncMock() - def _fake_run_coroutine_threadsafe(coro, loop): - coro.close() - return mock_future + other_state = bigquery_agent_analytics_plugin._LoopState( + write_client=mock_other_write_client, + batch_processor=mock_other_bp, + ) + plugin._loop_state_by_loop[other_loop] = other_state - with mock.patch.object( - asyncio, - "run_coroutine_threadsafe", - side_effect=_fake_run_coroutine_threadsafe, - ) as mock_rcts: - await plugin.shutdown() + await plugin.shutdown(timeout=5) - # Verify run_coroutine_threadsafe was called with - # the other loop. - mock_rcts.assert_called() - call_args = mock_rcts.call_args - assert call_args[0][1] is other_loop + # The drain ran on the OTHER loop's thread and the state was + # claimed after a clean completion. + assert drain_thread_ids == [thread.ident] + assert other_loop not in plugin._loop_state_by_loop + mock_other_write_client.transport.close.assert_awaited() + finally: + other_loop.call_soon_threadsafe(other_loop.stop) + thread.join(timeout=5) + other_loop.close() class TestCacheMetadataLogging: @@ -8364,6 +8479,105 @@ async def test_no_offloader_falls_back_to_truncate(self): assert parts[0]["storage_mode"] == "INLINE" assert "TRUNCATED" in parts[0]["text"] + @pytest.mark.asyncio + async def test_raw_prompt_text_is_sanitized_inline(self): + """Prompt, role, and system strings are redacted before row storage.""" + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + secret = "INLINE-CONTENT-SECRET" + request = llm_request_lib.LlmRequest( + contents=[ + types.Content( + role=json.dumps({"authorization": secret}), + parts=[types.Part(text=json.dumps({"secret": secret}))], + ) + ], + config=types.GenerateContentConfig( + system_instruction=json.dumps({"private_key": secret}) + ), + ) + + payload, parts, _ = await parser.parse(request) + stored = json.dumps({"content": payload, "content_parts": parts}) + assert secret not in stored + assert stored.count("[REDACTED]") >= 3 + + @pytest.mark.asyncio + async def test_raw_prompt_text_is_redacted_at_row_boundary( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + """The serialized BigQuery row never regains parser-redacted text.""" + secret = "ROW-BOUNDARY-CONTENT-SECRET" + request = llm_request_lib.LlmRequest( + model="gemini-pro", + contents=[ + types.Content( + role="user", + parts=[types.Part(text=json.dumps({"access_token": secret}))], + ) + ], + ) + + await bq_plugin_inst._log_event( + "LLM_REQUEST", + callback_context, + raw_content=request, + event_data=bigquery_agent_analytics_plugin.EventData( + model=request.model + ), + ) + await bq_plugin_inst.flush() + row = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + stored = json.dumps(row, default=str) + assert secret not in stored + assert "[REDACTED]" in stored + + @pytest.mark.asyncio + async def test_gcs_text_upload_receives_only_sanitized_content(self): + """Raw text is sanitized before either its GCS or row representation.""" + mock_offloader = mock.AsyncMock() + mock_offloader.upload_content.return_value = "gs://bucket/safe.txt" + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=mock_offloader, + trace_id="t", + span_id="s", + max_length=-1, + ) + secret = "GCS-CONTENT-SECRET" + text = json.dumps({"token": secret, "padding": "x" * (33 * 1024)}) + + payload, parts, _ = await parser.parse( + types.Content(parts=[types.Part(text=text)]) + ) + + uploaded = mock_offloader.upload_content.call_args.args[0] + assert secret not in uploaded + assert "[REDACTED]" in uploaded + stored = json.dumps({"content": payload, "content_parts": parts}) + assert secret not in stored + assert "[REDACTED]" in stored + + @pytest.mark.asyncio + async def test_internal_formatter_sentinel_is_preserved(self): + """Raw-text sanitization never corrupts generated formatter sentinels.""" + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, trace_id="t", span_id="s" + ) + payload, _, _ = await parser.parse( + bigquery_agent_analytics_plugin._FORMATTER_FAILED_SENTINEL + ) + assert payload == bigquery_agent_analytics_plugin._FORMATTER_FAILED_SENTINEL + # ================================================================ # TEST CLASS: AGENT_RESPONSE logging (Issue #87) @@ -8689,7 +8903,9 @@ async def fake_append_rows(requests, **kwargs): assert bp.dropped_event_count == 2 @pytest.mark.asyncio - async def test_non_retryable_drops_are_counted(self, dummy_arrow_schema): + async def test_non_retryable_drops_are_counted( + self, dummy_arrow_schema, caplog + ): bp = self._make_processor(dummy_arrow_schema) self._stub_arrow_prep(bp) @@ -8704,10 +8920,17 @@ async def fake_append_rows(requests, **kwargs): bp.write_client.append_rows.side_effect = fake_append_rows - await bp._write_rows_with_retry([{"a": 1}]) + secret = "NONRETRYABLE-ROW-SECRET" + with caplog.at_level( + logging.ERROR, + logger="google_adk.google.adk.plugins.bigquery_agent_analytics_plugin", + ): + await bp._write_rows_with_retry([{"a": secret}]) assert bp.get_drop_stats()["non_retryable"] == 1 assert bp.dropped_event_count == 1 + assert secret not in caplog.text + assert "1 row(s) dropped" in caplog.text def test_plugin_get_drop_stats_aggregates_across_loops( self, dummy_arrow_schema @@ -9843,8 +10066,8 @@ async def test_both_payload_columns_denied_skips_parse_and_offload( mock_blob.upload_from_string.assert_not_called() -class TestHardening: - """Safety and lifecycle invariants.""" +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 @@ -9886,7 +10109,7 @@ async def test_final_attributes_pass_redacts_direct_producers( These producers copy values into attributes without going through _recursive_smart_truncate; the final pre-serialization pass must - redact them. + redact them (#6356 P1-2). """ _ = mock_auth_default, mock_bq_client config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( @@ -9935,7 +10158,7 @@ async def test_final_attributes_pass_redacts_direct_producers( async def test_concurrent_parses_never_share_gcs_paths(self): """Two overlapping two-part parses keep call-local trace/span paths. - Regression: with identity stored on the shared parser, + 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. """ @@ -9982,9 +10205,7 @@ async def test_setup_failure_keeps_not_started_then_retries( mock_asyncio_to_thread, ): """Failed table readiness leaves _started=False, counts the loss, and - - a later event retries successfully. - """ + a later event retries successfully (#6356 P1-4).""" _ = mock_auth_default mock_bq_client.get_table.side_effect = cloud_exceptions.InternalServerError( "control plane hiccup" @@ -10011,7 +10232,7 @@ async def test_setup_failure_keeps_not_started_then_retries( 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. + # without ever re-checking the table (#6360 review P1-1). assert mock_bq_client.get_table.call_count == failed_calls + 1 @pytest.mark.asyncio @@ -10019,9 +10240,7 @@ 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. - """ + 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 @@ -10056,8 +10275,8 @@ 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. Both must still be - redacted; innocent strings stay unchanged. + 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 @@ -10084,7 +10303,7 @@ 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. + paths need the per-parse uid + content ordinal (#6360 review P1-2). """ uploaded: list[str] = [] @@ -10128,7 +10347,7 @@ async def test_formatter_failure_log_does_not_leak_payload( """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; only the exception + through exc_info tracebacks (#6360 review P1-3); only the exception class is logged. """ _ = mock_auth_default, mock_bq_client @@ -10152,7 +10371,10 @@ def leaky_formatter(content, event_type): ), ) assert "TOPSECRET-PAYLOAD" not in caplog.text - assert "ValueError" in caplog.text + # Round-10 P2-7: the message is CONSTANT — even the exception class + # name can be payload-derived, so it is no longer logged. + assert "Content formatter failed" in caplog.text + assert "ValueError" not in caplog.text @pytest.mark.asyncio async def test_shutdown_folds_processor_drops_into_stats( @@ -10161,8 +10383,7 @@ async def test_shutdown_folds_processor_drops_into_stats( """Processor drop counters survive shutdown via the plugin counters. get_drop_stats() used to read only live loop states, which shutdown() - clears. - """ + 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 @@ -10183,10 +10404,9 @@ async def test_shutdown_folds_processor_drops_into_stats( def test_setstate_backfills_new_runtime_fields( self, mock_auth_default, mock_bq_client ): - """Pickles from older code lack the new fields; __setstate__ must - - backfill them so get_drop_stats()/_ensure_started don't raise. - """ + """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 @@ -10207,7 +10427,8 @@ def test_setstate_backfills_new_runtime_fields( def test_invalid_config_rejects_nan_and_wrong_types( self, mock_auth_default, mock_bq_client ): - """NaN and wrong-typed values must fail construction: ordered comparisons alone let NaN pass every range check.""" + """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") @@ -10234,9 +10455,9 @@ def test_invalid_config_rejects_nan_and_wrong_types( 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 - . + 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]"}' @@ -10248,8 +10469,10 @@ 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. + 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"}) @@ -10262,23 +10485,31 @@ def test_mapping_views_are_redacted(self): assert out["userdict"]["refresh_token"] == "[REDACTED]" def test_deep_json_blob_fails_closed(self): - """A blob too deep to inspect becomes a sentinel, not a pass-through. + """A blob beyond the fixed nesting limit fails closed on every runtime. - Structural nesting beyond the sanitizer's depth bound cannot be verified - secret-free, so it fails closed regardless of the interpreter's json - recursion handling; the row keeps flowing with the blob replaced. + Older Python runtimes raise RecursionError while Python 3.14's iterative + JSON decoder accepts this input. The row must keep flowing with the same + whole-blob sentinel regardless (#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]" + def test_json_nesting_limit_ignores_brackets_inside_strings(self): + """Payload punctuation does not count as structural JSON nesting.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + blob = json.dumps({"note": "prose " + "[" * 1001 + "]" * 1001}) + out, truncated = truncate({"blob": blob}, 500 * 1024) + assert out["blob"] == blob + assert truncated is False + @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(). + rows by the drain in shutdown() (#6360 review round 2 P2-4). """ write_started = asyncio.Event() @@ -10313,9 +10544,7 @@ 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 @@ -10335,9 +10564,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. - """ + 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 @@ -10358,7 +10585,7 @@ async def test_gcs_uploads_use_full_uid_and_create_only(self): 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. + to another event's bytes (#6360 review round 2 P2-8). """ uploaded: list[str] = [] @@ -10398,7 +10625,8 @@ def test_unmaterializable_json_blob_fails_closed(self): 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. + 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 + "}" @@ -10411,7 +10639,7 @@ def test_label_only_upgrade_failure_does_not_block_readiness(self): The table schema is write-compatible; only the governance label is stale. Blocking readiness turned every event into setup_unavailable - although writes would succeed. + although writes would succeed (#6360 review round 3 P1-2). """ config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( auto_schema_upgrade=True, @@ -10443,10 +10671,12 @@ def test_ensure_started_coalesces_across_event_loops( """_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. Per-loop locks make - each loop coalesce independently. + '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 ) @@ -10482,7 +10712,8 @@ def test_concurrent_stale_cleanup_folds_once( """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. + KeyError; the pop-claim makes folding idempotent (#6360 review round 3 + P2). """ _ = mock_auth_default, mock_bq_client plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( @@ -10502,9 +10733,7 @@ def test_concurrent_stale_cleanup_folds_once( @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): @@ -10537,7 +10766,7 @@ 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. + redaction, including with escaped keys (#6360 review round 4 P1-1). """ truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate cases = [ @@ -10554,7 +10783,7 @@ 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. + 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 + '"}' @@ -10570,12 +10799,14 @@ def test_over_limit_blob_never_parsed(self): def test_shared_setup_runs_exactly_once_across_loops( self, mock_auth_default, mock_bq_client ): - """Concurrent loops coalesce onto ONE shared setup. + """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 ) @@ -10605,7 +10836,7 @@ def run_in_fresh_loop(): for t in threads: t.start() # Deterministic rendezvous: hold the owner inside setup until BOTH - # threads have entered _ensure_started. + # threads have entered _ensure_started (#6360 review round 5 P2). entered.wait(timeout=5) release.set() for t in threads: @@ -10623,6 +10854,8 @@ def test_failed_shared_setup_is_consistent_across_loops( ): """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 ) @@ -10677,12 +10910,16 @@ async def test_namedtuple_attribute_does_not_drop_row( dummy_arrow_schema, mock_asyncio_to_thread, ): - """A namedtuple in attributes serializes as a list, not a TypeError. + """A namedtuple in attributes serializes as a mapping, not a TypeError. Reconstructing tuple subclasses positionally raised in the final pass - and the safe callback dropped the entire row. + and the safe callback dropped the entire row (#6360 review round 4 P2); + round 7 P1-2 then required the mapping shape so field-name redaction + can run. """ _ = 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 @@ -10702,59 +10939,244 @@ async def test_namedtuple_attribute_does_not_drop_row( mock_write_client, dummy_arrow_schema ) attrs = json.loads(log_entry["attributes"]) - assert attrs["point"] == [1, 2] + # Namedtuples serialize as a MAPPING so field-name redaction can run + # (#6360 round 7 P1-2 superseded the round-4 plain-list shape); the + # row is still emitted either way. + assert attrs["point"] == {"x": 1, "y": 2} - def test_setup_future_leaves_no_loop_references( + @pytest.mark.asyncio + async def test_setup_blocked_before_loop_state_does_not_leak( 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 - ; the cross-loop future replaces it. - """ + """Round-7 P1-5: a shutdown() that completes while setup is blocked + creating the shared client must abort the resumed setup, publish + nothing, and release every resource the attempt created.""" _ = mock_auth_default, mock_bq_client plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( PROJECT_ID, DATASET_ID, table_id=TABLE_ID ) + plugin._credentials = mock.MagicMock(quota_project_id=None) - async def noop_setup(**kwargs): - return None + entered = threading.Event() + release = threading.Event() - 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 gated_client(*args, **kwargs): + del args, kwargs + entered.set() + release.wait(10) + return mock.MagicMock() - def test_cleanup_survives_concurrent_insertion( + with mock.patch( + "google.adk.plugins.bigquery_agent_analytics_plugin.bigquery.Client", + side_effect=gated_client, + ): + owner = asyncio.create_task(plugin._ensure_started()) + while not entered.is_set(): + await asyncio.sleep(0.01) + # Shutdown completes fully while setup is blocked in the executor. + await plugin.shutdown() + release.set() + outcome = await owner # aborts internally; never raises + + assert outcome == "aborted" + assert plugin._started is False + assert plugin.client is None + assert plugin._executor is None + assert plugin.parser is None + assert plugin.offloader is None + assert plugin._loop_state_by_loop == {} + # Round-8 P2-9: a direct start (no row) records no phantom loss; the + # structured outcome lets the row owner count instead. + assert plugin.get_drop_stats().get("shutdown_race", 0) == 0 + # The abort is not a service failure: no poisoned backoff window. + assert plugin._startup_error is None + + def test_concurrent_shutdown_folds_counters_once( self, mock_auth_default, mock_bq_client ): - """Cleanup snapshots keys, so insertion during is_closed() cannot raise + """Round-7 P1-6: two threads racing into shutdown() must not both be + admitted — the same processors' drop counters were folded twice.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) - 'dictionary changed size during iteration'. - """ + def make_state(): + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock() + state.batch_processor.get_drop_stats = mock.MagicMock( + return_value={"queue_full": 1} + ) + return state + + for _ in range(2): + # Closed fakes: shutdown claims and folds them without scheduling + # coroutines on them (a non-closed MagicMock loop leaked unawaited + # AsyncMock coroutines — #6360 round 8 verification note). + fake_loop = mock.MagicMock(spec=asyncio.AbstractEventLoop) + fake_loop.is_closed.return_value = True + plugin._loop_state_by_loop[fake_loop] = make_state() + + barrier = threading.Barrier(2) + errors = [] + + def run_shutdown(): + try: + barrier.wait(timeout=10) + asyncio.run(plugin.shutdown(timeout=0.1)) + except Exception as e: # pylint: disable=broad-except + errors.append(e) + + threads = [threading.Thread(target=run_shutdown) for _ in range(2)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=30) + assert not t.is_alive() + + assert not errors + # Two states, one queue_full each: exactly one shutdown owner folds + # them, so anything above 2 means double-folding. + assert plugin.get_drop_stats().get("queue_full", 0) == 2 + + def test_drop_counters_are_thread_safe( + self, mock_auth_default, mock_bq_client + ): + """Round-7 P2-7: concurrent _count_local_drop() increments from + multiple threads must not lose updates.""" _ = 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} + increments = 5_000 + n_threads = 4 + old_interval = sys.getswitchinterval() + sys.setswitchinterval(1e-5) + try: - def is_closed_and_mutate(): - # Simulates another thread inserting mid-scan. - plugin._loop_state_by_loop[mock.MagicMock()] = mock.MagicMock() - return True + def worker(): + for _ in range(increments): + plugin._count_local_drop("stress") - dead_loop.is_closed.side_effect = is_closed_and_mutate - plugin._loop_state_by_loop[dead_loop] = state + threads = [threading.Thread(target=worker) for _ in range(n_threads)] + for t in threads: + t.start() + for t in threads: + t.join(timeout=60) + assert not t.is_alive() + finally: + sys.setswitchinterval(old_interval) - plugin._cleanup_stale_loop_states() # must not raise - assert plugin.get_drop_stats().get("write_failed") == 7 + assert plugin.get_drop_stats()["stress"] == increments * n_threads + + def test_unlimited_mode_scans_entire_emitted_value(self): + """Round-15 P1-1: in unlimited mode the ENTIRE emitted value is + classified — a credential document just past the inspection window + fails closed; escape-free giant quoted prose passes whole.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + ceiling = bigquery_agent_analytics_plugin._MAX_JSON_INSPECT_CHARS + + raw = ( + '"' + + "a" * (ceiling + 10) + + '\\u007b\\"access\\u005ftoken\\":' + + '\\"R15-UNLIMITED-SECRET\\"\\u007d"' + ) + out, truncated = truncate({"blob": raw}, -1) + assert "R15-UNLIMITED-SECRET" not in json.dumps(out) + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True + + prose = '"' + "hello world " * ((ceiling // 12) + 10) + '"' + out, truncated = truncate({"s": prose}, -1) + assert out["s"] == prose + assert truncated is False + + def test_normalizer_redacts_and_bounds(self): + """Round-15 P1-2/P1-3: the JSON-native normalizer applies + sensitive-key/temp: redaction and the configured length bound, while + preserving sentinels and bracketed prose.""" + normalize = bigquery_agent_analytics_plugin._normalize_json_native + + out, _ = normalize( + {"prompt": [{"role": {"access_token": "R15-NATIVE-SECRET"}}]}, + 10000, + ) + assert "R15-NATIVE-SECRET" not in json.dumps(out) + assert out["prompt"][0]["role"]["access_token"] == "[REDACTED]" + + out, replaced = normalize("R15-ROLE-" + "x" * 1_000_000, 10) + assert out == "R15-ROLE-x...[TRUNCATED]" + assert replaced is True + + for preserved in ("[FORMATTER_FAILED]", "[bracketed] prose"): + out, replaced = normalize(preserved, 10000) + assert out == preserved + assert replaced is False + + def test_normalizer_preserves_post_normalization_key_collisions(self): + """Normalized keys never silently overwrite an earlier value.""" + normalize = bigquery_agent_analytics_plugin._normalize_json_native + + unsupported_first, replaced = normalize( + {object(): "unsupported", "[UNSUPPORTED_KEY_1]": "genuine"}, + 10000, + ) + assert unsupported_first == { + "[UNSUPPORTED_KEY_1]": "unsupported", + "[KEY_COLLISION_2][UNSUPPORTED_KEY_1]": "genuine", + } + assert replaced is True + + genuine_first, replaced = normalize( + {"[UNSUPPORTED_KEY_1]": "genuine", object(): "unsupported"}, + 10000, + ) + assert genuine_first == { + "[UNSUPPORTED_KEY_1]": "genuine", + "[KEY_COLLISION_2][UNSUPPORTED_KEY_1]": "unsupported", + } + assert replaced is True + + marker_reserved, replaced = normalize( + { + object(): "unsupported", + "[KEY_COLLISION_2][UNSUPPORTED_KEY_1]": "reserved", + "[UNSUPPORTED_KEY_1]": "genuine", + }, + 10000, + ) + assert marker_reserved == { + "[UNSUPPORTED_KEY_1]": "unsupported", + "[KEY_COLLISION_2][UNSUPPORTED_KEY_1]": "reserved", + "[KEY_COLLISION_3][UNSUPPORTED_KEY_1]": "genuine", + } + assert replaced is True + + budget_reserved, replaced = normalize( + { + "[SANITIZE_BUDGET_EXCEEDED]": "reserved", + "[KEY_COLLISION_1][SANITIZE_BUDGET_EXCEEDED]": "marker", + "omitted": "value", + }, + 10000, + budget=[3], + ) + assert budget_reserved == { + "[SANITIZE_BUDGET_EXCEEDED]": "reserved", + "[KEY_COLLISION_1][SANITIZE_BUDGET_EXCEEDED]": "marker", + "[KEY_COLLISION_2][SANITIZE_BUDGET_EXCEEDED]": ( + "[SANITIZE_BUDGET_EXCEEDED]" + ), + } + assert replaced is True @pytest.mark.asyncio - async def test_depth_capped_payload_flags_row_truncated( + async def test_native_secret_mapping_via_model_field_redacted( self, mock_write_client, invocation_context, @@ -10765,255 +11187,3038 @@ async def test_depth_capped_payload_flags_row_truncated( dummy_arrow_schema, mock_asyncio_to_thread, ): - """A real payload cut off by the depth cap marks the ROW as truncated - - . - """ + """Round-15 P1-2 at the row boundary: a nested model property handing + the parser a raw credential mapping is redacted in the written row.""" _ = mock_auth_default, mock_bq_client - deep: dict = {"leaf": "payload"} - for _ in range(60): - deep = {"level": deep} + + class EvilContent(types.Content): + + def __getattribute__(self, name): + if name == "role": + return {"access_token": "R15-NATIVE-SECRET"} + return super().__getattribute__(name) + + hostile = llm_request_lib.LlmRequest(contents=[EvilContent(parts=[])]) + assert type(hostile) is llm_request_lib.LlmRequest + + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=lambda content, event_type: hostile + ) async with managed_plugin( - PROJECT_ID, DATASET_ID, table_id=TABLE_ID + 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", + "LLM_REQUEST", callback_context, - event_data=bigquery_agent_analytics_plugin.EventData( - extra_attributes={"deep": deep}, - ), + event_data=bigquery_agent_analytics_plugin.EventData(), ) 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 + assert "R15-NATIVE-SECRET" not in json.dumps(log_entry, default=str) + assert "[REDACTED]" in json.dumps(log_entry, default=str) - 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 + def test_overlimit_prose_prefixed_encoded_string_fails_closed(self): + """Round-14 P1-1: an over-limit quoted value whose EMITTED prefix + hides an escaped container after prose fails closed; escape-free + over-limit quoted prose still raw-truncates.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate - . - """ - _ = 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 - ) + raw = ( + '"note \\u007b\\"access\\u005ftoken\\":' + '\\"R14-OVERLIMIT-SECRET\\"\\u007d' + + "x" * 10050 + + '"' ) - plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( - PROJECT_ID, DATASET_ID, table_id=TABLE_ID, config=config - ) - assert plugin.config.retry_config.max_retries == 0 + out, truncated = truncate({"blob": raw}, 10000) + assert "R14-OVERLIMIT-SECRET" not in json.dumps(out) + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True + + prose = '"' + "hello world " * 2000 + '"' + out, truncated = truncate({"s": prose}, 1000) + assert out["s"].endswith("...[TRUNCATED]") + assert truncated is True @pytest.mark.asyncio - async def test_owner_cancellation_does_not_poison_rendezvous( - self, mock_auth_default, mock_bq_client + async def test_late_detonating_nested_model_normalized( + 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, + caplog, ): - """A cancelled setup owner finalizes the shared future so later - - startups are not stuck forever. - """ + """Round-14 P1-2: a nested model that parses cleanly but plants an + object whose __repr__ raises must be normalized inside the parse + boundary — the row survives Arrow preparation and the payload never + reaches the logs.""" _ = 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) + class LateBomb: - 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 + def __repr__(self): + raise RuntimeError("R14-LATE-SERIALIZE-SECRET") - assert plugin._setup_future is None # rendezvous cleared + class EvilContent(types.Content): - # A later attempt is not stuck: it claims a fresh future and runs. - async def ok_setup(**kwargs): - return None + def __getattribute__(self, name): + if name == "role": + return LateBomb() + return super().__getattribute__(name) - 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 + hostile = llm_request_lib.LlmRequest(contents=[EvilContent(parts=[])]) + assert type(hostile) is llm_request_lib.LlmRequest + + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=lambda content, event_type: hostile + ) + 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) + with caplog.at_level(logging.WARNING): + await plugin._log_event( + "LLM_REQUEST", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + await asyncio.sleep(0.01) + # The row survives Arrow preparation (exercised by the capture + # helper) with the hostile object replaced by a sentinel. + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + dumped = json.dumps(log_entry, default=str) + assert "R14-LATE-SERIALIZE-SECRET" not in dumped + assert "R14-LATE-SERIALIZE-SECRET" not in caplog.text + assert "[UNSUPPORTED_OBJECT]" in dumped + assert plugin.get_drop_stats().get("arrow_prep_failed", 0) == 0 @pytest.mark.asyncio - async def test_waiter_cancellation_does_not_cancel_shared_future( + async def test_remote_scheduling_failure_keeps_teardown_incomplete( self, mock_auth_default, mock_bq_client ): - """Cancelling one waiter must not cancel the owner's shared future - - . - """ + """Round-14 P1-3: an exception from _schedule_remote_drain() itself + counts the state as retained, so shutdown raises instead of + reporting success over live state.""" _ = 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() + remote_loop = asyncio.new_event_loop() + thread = threading.Thread(target=remote_loop.run_forever, daemon=True) + thread.start() + try: + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.get_drop_stats = mock.MagicMock(return_value={}) + plugin._loop_state_by_loop[remote_loop] = state - async def gated_setup(**kwargs): - entered.set() - await release.wait() + with mock.patch.object( + plugin, + "_schedule_remote_drain", + side_effect=RuntimeError("loop closed during scheduling"), + ): + with pytest.raises( + bigquery_agent_analytics_plugin._ShutdownIncompleteError + ): + await plugin.shutdown(timeout=2) + assert remote_loop in plugin._loop_state_by_loop + finally: + remote_loop.call_soon_threadsafe(remote_loop.stop) + thread.join(timeout=5) + remote_loop.close() + + def test_prose_inside_encoded_string_fails_closed(self): + """Round-13 P1-2: a single valid encoded string whose DECODED content + hides a container after prose fails closed; ordinary quoted prose + (including inner quotes) passes through.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate - 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 + raw = '"note \\u007b\\"access\\u005ftoken\\":\\"R13-SECRET\\"\\u007d"' + out, truncated = truncate({"blob": raw}, 10000) + assert "R13-SECRET" not in json.dumps(out) + assert out["blob"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True - assert plugin._started is True + for prose in ( + '"just quoted prose"', + json.dumps('he said "hi"'), + ): + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False @pytest.mark.asyncio - async def test_shutdown_wins_over_in_flight_setup( - self, mock_auth_default, mock_bq_client + async def test_nested_hostile_model_subclass_fails_closed( + 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, + caplog, ): - """Setup completing after shutdown() must not resurrect _started - - . - """ + """Round-13 P1-1: a hostile model subclass NESTED inside an + exact-typed formatter result fails closed at the parse boundary — the + row is written with a sentinel and the payload-bearing exception + never reaches the logs.""" _ = 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() + class EvilPart(types.Part): - 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 + def __getattribute__(self, name): + if name == "file_data": + raise RuntimeError("R13-NESTED-SECRET") + return super().__getattribute__(name) - assert plugin._started is False - assert plugin.get_drop_stats().get("shutdown_race", 0) >= 1 + hostile = types.Content(parts=[EvilPart()]) + assert type(hostile) is types.Content # passes the exact-type gate + + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=lambda content, event_type: hostile + ) + 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) + with caplog.at_level(logging.WARNING): + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert "R13-NESTED-SECRET" not in json.dumps(log_entry, default=str) + assert "R13-NESTED-SECRET" not in caplog.text + assert "[CONTENT_PARSE_FAILED]" in log_entry["content"] + assert log_entry["is_truncated"] is True + assert plugin.get_drop_stats().get("content_parse_failed", 0) == 1 @pytest.mark.asyncio - async def test_close_invokes_full_shutdown( + async def test_failed_remote_drain_fails_coalesced_waiter_too( self, mock_auth_default, mock_bq_client ): - """plugin.close() (Runner/PluginManager ownership) performs the real - - shutdown instead of the inherited no-op. - """ + """Round-13 P1-3: a failed remote drain keeps teardown incomplete for + the coalesced waiter as well — neither caller reports success over + live state.""" _ = 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): - """Additional blob shapes: bytes/bytearray blobs, BOM-prefixed JSON, + remote_loop = asyncio.new_event_loop() + thread = threading.Thread(target=remote_loop.run_forever, daemon=True) + thread.start() + try: + entered = threading.Event() + release = threading.Event() + + async def failing_drain(timeout=None): + del timeout + entered.set() + while not release.is_set(): + await asyncio.sleep(0.01) + raise RuntimeError("remote drain fails") + + bp = mock.MagicMock(spec=bigquery_agent_analytics_plugin.BatchProcessor) + bp.shutdown = failing_drain + bp.get_drop_stats = mock.MagicMock(return_value={}) + state = mock.MagicMock() + state.write_client = None + state.batch_processor = bp + plugin._loop_state_by_loop[remote_loop] = state + + owner = asyncio.create_task(plugin.shutdown(timeout=5)) + while not entered.is_set(): + await asyncio.sleep(0.01) + waiter = asyncio.create_task(plugin.shutdown()) + await asyncio.sleep(0.05) + release.set() + with pytest.raises( + bigquery_agent_analytics_plugin._ShutdownIncompleteError + ): + await owner + # The retrying waiter hits the same persistent remote failure. + with pytest.raises( + bigquery_agent_analytics_plugin._ShutdownIncompleteError + ): + await asyncio.wait_for(waiter, timeout=10) + assert remote_loop in plugin._loop_state_by_loop + finally: + remote_loop.call_soon_threadsafe(remote_loop.stop) + thread.join(timeout=5) + remote_loop.close() - __str__-returned credential JSON, and Mapping converter results. - """ + def test_prose_then_encoded_document_redacted(self): + """Round-12 P1-1: an encoded credential document after a stretch of + raw prose in the suffix is still decoded and redacted.""" truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate - class ToDictMapping: + v = '"note" then "\\u007b\\"access_token\\":\\"R12-SECRET\\"\\u007d"' + out, truncated = truncate({"b": v}, 10000) + assert "R12-SECRET" not in json.dumps(out) + assert "[REDACTED]" in out["b"] + del truncated - def to_dict(self): - return collections.UserDict({"access_token": "SECRET-MAPPING"}) + # A chain of prose and documents is walked to the depth cap. + chain = ( + '"note" one "plain" two' + ' "\\u007b\\"refresh_token\\":\\"R12-CHAIN-SECRET\\"\\u007d"' + ) + out, _ = truncate({"b": chain}, 10000) + assert "R12-CHAIN-SECRET" not in json.dumps(out) - class StrLeaker: + # An escape hidden in a prose gap cannot be verified. + out, truncated = truncate({"s": '"note" \\then "x"'}, 10000) + assert out["s"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True - def __str__(self): - return '{"access_token": "SECRET-STR"}' + # Multi-quote prose still passes through. + prose = '"a" and then "b" happened' + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False - 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 + @pytest.mark.asyncio + async def test_native_subclass_formatter_result_fails_closed( + 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, + caplog, + ): + """Round-12 P1-2: a SUBCLASS of a parser-native model shape from the + formatter fails closed at the boundary instead of reaching parser + attribute accesses outside it.""" + _ = mock_auth_default, mock_bq_client - def test_sanitizer_stops_at_node_budget(self): - """A very wide payload stops at the work budget and flags truncation + class SubRequest(llm_request_lib.LlmRequest): + pass - . - """ - 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 + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=lambda content, event_type: SubRequest() ) - assert len(out["wide"]) <= len(wide) + 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) + with caplog.at_level(logging.WARNING): + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert ( + bigquery_agent_analytics_plugin._FORMATTER_FAILED_SENTINEL + in log_entry["content"] + ) + assert "SubRequest" not in caplog.text + assert plugin.get_drop_stats().get("formatter_failed", 0) == 1 @pytest.mark.asyncio - async def test_stale_loop_cleanup_counts_queued_rows( + async def test_persistent_teardown_failure_raises_to_all_callers( self, mock_auth_default, mock_bq_client ): - """Queued rows on a closed loop are counted under stale_loop - - . - """ + """Round-12 P1-3: a persistently failing teardown must not report + success to the owner or to retrying waiters.""" _ = 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 + entered = asyncio.Event() + release = asyncio.Event() + calls = [] + + async def always_failing_drain(timeout=None): + del timeout + calls.append(1) + if len(calls) == 1: + entered.set() + await release.wait() + raise RuntimeError("drain always fails") + 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 + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock( + side_effect=always_failing_drain + ) + state.batch_processor.get_drop_stats = mock.MagicMock(return_value={}) + plugin._loop_state_by_loop[asyncio.get_running_loop()] = state - plugin._cleanup_stale_loop_states() - assert plugin.get_drop_stats().get("stale_loop") == 1 + owner = asyncio.create_task(plugin.shutdown(timeout=5)) + await entered.wait() + waiter = asyncio.create_task(plugin.shutdown()) + await asyncio.sleep(0.05) + release.set() + with pytest.raises(RuntimeError, match="drain always fails"): + await owner + # The retrying waiter becomes the owner, fails the same way, and + # surfaces the failure instead of returning success over live state. + with pytest.raises(RuntimeError, match="drain always fails"): + await asyncio.wait_for(waiter, timeout=5) + assert len(calls) == 2 + assert plugin._loop_state_by_loop != {} + + def test_unicode_escaped_trailing_document_redacted(self): + """Round-11 P1-1: a trailing quoted JSON document whose decoded + content hides a container behind Unicode escapes is decoded and + redacted; quoted prose and prose suffixes stay untouched.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + v = '"note" "\\u007b\\"access_token\\":\\"R11-SECRET\\"\\u007d"' + out, truncated = truncate({"b": v}, 10000) + assert "R11-SECRET" not in json.dumps(out) + assert "[REDACTED]" in out["b"] + del truncated + + # A stray leading escape in the suffix cannot be classified. + out, truncated = truncate({"s": '"note" \\x'}, 10000) + assert out["s"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True + + for prose in ( + '"hello" she said', + '"a" and then "b" happened', + '"note" "just more prose"', + ): + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False @pytest.mark.asyncio - async def test_restart_rebuilds_parser_and_offloader( + async def test_waiter_retries_after_failed_owner_teardown( self, mock_auth_default, mock_bq_client ): - """shutdown() clears parser/offloader so a restart cannot reuse the - - terminated executor. - """ + """Round-11 P1-3: an ordinary teardown exception must not report + successful completion to coalesced waiters; they retry ownership.""" _ = 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 + entered = asyncio.Event() + release = asyncio.Event() + calls = [] + + async def failing_first_drain(timeout=None): + del timeout + calls.append(1) + if len(calls) == 1: + entered.set() + await release.wait() + raise RuntimeError("first drain fails") + + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock( + side_effect=failing_first_drain + ) + state.batch_processor.get_drop_stats = mock.MagicMock(return_value={}) + plugin._loop_state_by_loop[asyncio.get_running_loop()] = state + + owner = asyncio.create_task(plugin.shutdown(timeout=5)) + await entered.wait() + waiter = asyncio.create_task(plugin.shutdown()) + await asyncio.sleep(0.05) + release.set() + # Round-12 P1-3: the OWNER must not report success over live state — + # the teardown error propagates to its caller. + with pytest.raises(RuntimeError, match="first drain fails"): + await owner + # The waiter must not have accepted the failed teardown as success: + # it retries ownership, the second drain succeeds, state is claimed. + await asyncio.wait_for(waiter, timeout=5) + assert len(calls) == 2 + assert plugin._loop_state_by_loop == {} + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("error::RuntimeWarning") + async def test_rejecting_task_factory_does_not_leak_coroutine( + self, mock_auth_default, mock_bq_client + ): + """Round-11 P2-4: if the remote loop's task factory rejects task + creation, the drain coroutine is closed instead of leaking.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + remote_loop = asyncio.new_event_loop() + + def rejecting_factory(loop, coro, **kwargs): + del loop, coro, kwargs + raise RuntimeError("factory rejects") + + remote_loop.set_task_factory(rejecting_factory) + thread = threading.Thread(target=remote_loop.run_forever, daemon=True) + thread.start() + try: + state = mock.MagicMock() + state.write_client = None + bp = mock.MagicMock(spec=bigquery_agent_analytics_plugin.BatchProcessor) + + async def drain(timeout=None): + del timeout + + bp.shutdown = drain + bp.get_drop_stats = mock.MagicMock(return_value={}) + state.batch_processor = bp + plugin._loop_state_by_loop[remote_loop] = state + + # Round-13 P1-3: the failed drain keeps teardown incomplete. + with pytest.raises( + bigquery_agent_analytics_plugin._ShutdownIncompleteError + ): + await plugin.shutdown(timeout=2) + # The failed drain retains the state; no never-awaited warning + # (filterwarnings turns it into a hard error). + assert remote_loop in plugin._loop_state_by_loop + finally: + remote_loop.call_soon_threadsafe(remote_loop.stop) + thread.join(timeout=5) + remote_loop.close() + + def test_unterminated_quoted_container_fails_closed(self): + """Round-10 P1-1: a quoted layer that visibly begins an encoded + container but is missing its final quote fails closed; unterminated + quoted prose passes through.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + v = '"{\\"access_token\\":\\"R10-SECRET\\"}' + out, truncated = truncate({"cache": v}, 10000) + assert "R10-SECRET" not in json.dumps(out) + assert out["cache"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True + + prose = '"unterminated prose without a container' + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False + + @pytest.mark.asyncio + async def test_identity_formatter_preserves_native_shapes( + 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, + ): + """Round-10 P1-2: an identity formatter must not destroy parser-native + shapes (dict/list) that it returns untransformed.""" + _ = mock_auth_default, mock_bq_client + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=lambda content, event_type: content + ) + 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, + raw_content={"response": "safe-dict-content"}, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert "safe-dict-content" in json.dumps(log_entry, default=str) + assert ( + bigquery_agent_analytics_plugin._FORMATTER_FAILED_SENTINEL + not in json.dumps(log_entry, default=str) + ) + assert plugin.get_drop_stats().get("formatter_failed", 0) == 0 + + @pytest.mark.asyncio + async def test_waiter_retries_after_cancelled_owner_shutdown( + self, mock_auth_default, mock_bq_client + ): + """Round-10 P1-3: a coalesced caller must not claim success when the + owning shutdown was cancelled mid-teardown; it retries ownership and + finishes the job.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + gate = asyncio.Event() + calls = [] + + async def gated_first_shutdown(timeout=None): + del timeout + calls.append(1) + if len(calls) == 1: + await gate.wait() + + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock( + side_effect=gated_first_shutdown + ) + state.batch_processor.get_drop_stats = mock.MagicMock(return_value={}) + plugin._loop_state_by_loop[asyncio.get_running_loop()] = state + + owner = asyncio.create_task(plugin.shutdown(timeout=5)) + await asyncio.sleep(0.05) + waiter = asyncio.create_task(plugin.shutdown()) + await asyncio.sleep(0.05) + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + + # The waiter retried ownership and completed the teardown. + await asyncio.wait_for(waiter, timeout=5) + assert plugin._loop_state_by_loop == {} + assert len(calls) == 2 + + @pytest.mark.asyncio + async def test_slow_remote_drain_is_retained_without_close_errors( + self, mock_auth_default, mock_bq_client + ): + """Round-10 P1-4: a remote drain still running at the deadline is + retained and keeps running remotely; the caller never closes a + coroutine it no longer owns (no 'coroutine already executing').""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + remote_loop = asyncio.new_event_loop() + thread = threading.Thread(target=remote_loop.run_forever, daemon=True) + thread.start() + try: + release = threading.Event() + drain_finished = threading.Event() + + async def slow_drain(timeout=None): + del timeout + while not release.is_set(): + await asyncio.sleep(0.01) + drain_finished.set() + + bp = mock.MagicMock(spec=bigquery_agent_analytics_plugin.BatchProcessor) + bp.shutdown = slow_drain + bp.get_drop_stats = mock.MagicMock(return_value={}) + state = mock.MagicMock() + state.write_client = None + state.batch_processor = bp + plugin._loop_state_by_loop[remote_loop] = state + + # Round-13 P1-3: the timed-out drain keeps teardown incomplete. + with pytest.raises( + bigquery_agent_analytics_plugin._ShutdownIncompleteError + ): + await plugin.shutdown(timeout=0.2) + # Timed out: state retained, no ValueError from closing a running + # coroutine (shutdown would have logged/raised through its guard). + assert remote_loop in plugin._loop_state_by_loop + # The remote drain keeps running to completion on its own loop. + release.set() + assert drain_finished.wait(timeout=5) + finally: + remote_loop.call_soon_threadsafe(remote_loop.stop) + thread.join(timeout=5) + remote_loop.close() + + @pytest.mark.asyncio + async def test_formatter_logs_never_carry_payload_derived_names( + 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, + caplog, + ): + """Round-10 P2-7: formatter log lines are constant — payload-derived + result/exception CLASS NAMES never reach the logs.""" + _ = mock_auth_default, mock_bq_client + secret_result_cls = type("R10_RESULT_SECRET", (), {}) + secret_error_cls = type("R10_ERROR_SECRET", (Exception,), {}) + + outcomes = iter([secret_result_cls(), None]) + + def formatter(content, event_type): + del content, event_type + value = next(outcomes) + if value is None: + raise secret_error_cls() + return value + + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=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): + for _ in range(2): + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + assert "R10_RESULT_SECRET" not in caplog.text + assert "R10_ERROR_SECRET" not in caplog.text + assert plugin.get_drop_stats().get("formatter_failed", 0) == 2 + + def test_unicode_ws_and_bom_quoted_layers_fail_closed(self): + """Round-8 P1-1: BOM/NBSP/EM-SPACE prefixes inside quoted JSON layers + are normalized before every shape check, under- and over-limit.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + for pad in ("\u00a0", "\u2003", "\ufeff"): + v = json.dumps(pad + json.dumps({"access_token": "R8-WS-SECRET"})) + out, _ = truncate({"b": v}, 10000) + assert "R8-WS-SECRET" not in json.dumps(out), repr(pad) + + secret = "R8-WS-OVER-SECRET-" + "x" * 300 + for pad in ("\u00a0", "\u2003", "\ufeff"): + v = json.dumps( + pad + json.dumps({"access_token": secret}), ensure_ascii=False + ) + out, truncated = truncate({"b": v}, 120) + assert "R8-WS-OVER-SECRET" not in json.dumps(out), repr(pad) + assert truncated + + def test_quoted_prefix_suffix_smuggling_fails_closed(self): + """Round-8 P1-2: credential JSON smuggled after a harmless quoted + prefix fails closed; container-free quoted prose passes through.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + for v in ( + '"note" {"access_token":"R8-SUFFIX-SECRET"}', + '"note" blah {"refresh_token":"R8-SUFFIX-SECRET-2"}', + ): + out, truncated = truncate({"b": v}, 10000) + assert "R8-SUFFIX-SECRET" not in json.dumps(out) + assert truncated + + prose = '"hello" she said, "twice"' + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False + + def test_safe_scalar_subclass_str_not_published(self): + """Round-8 P1-3: subclasses of allowlisted scalar types cannot leak + values through an overridden __str__; base conversions are used.""" + import enum + import pathlib + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + class Credential(enum.Enum): + access_token = "R8-ENUM-SECRET" + + def __str__(self): + return self.value + + out, _ = truncate({"c": Credential.access_token}, 10000) + assert "R8-ENUM-SECRET" not in json.dumps(out) + assert out["c"] == "Credential.access_token" + + class SneakyPath(pathlib.PurePosixPath): + + def __str__(self): + return "R8-PATH-SECRET" + + out, _ = truncate({"p": SneakyPath("/tmp/x")}, 10000) + assert "R8-PATH-SECRET" not in json.dumps(out) + assert out["p"] == "/tmp/x" + + def test_safe_scalar_truncation_reports_flag(self): + """Round-8 P2-11: an over-limit safe scalar reports truncation.""" + import pathlib + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + out, truncated = truncate(pathlib.PurePosixPath("x" * 40), 8) + assert "[TRUNCATED]" in out + assert truncated is True + + def test_hostile_container_protocols_fail_closed(self): + """Round-8 P1-4: raising items()/iteration/field access fails closed + to a sentinel instead of escaping the sanitizer.""" + import collections.abc + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + class EvilMapping(collections.abc.Mapping): + + def __getitem__(self, k): + raise KeyError(k) + + def __len__(self): + return 1 + + def __iter__(self): + return iter(["a"]) + + def items(self): + raise RuntimeError("R8-MAPPING-SECRET") + + class EvilList(list): + + def __iter__(self): + raise RuntimeError("R8-LIST-SECRET") + + out, truncated = truncate({"m": EvilMapping(), "l": EvilList([1])}, 10000) + assert out["m"] == "[UNSUPPORTED_OBJECT]" + assert out["l"] == "[UNSUPPORTED_OBJECT]" + assert truncated is True + assert "R8-MAPPING-SECRET" not in json.dumps(out) + + @pytest.mark.asyncio + async def test_hostile_protocol_row_still_emitted_no_canary_in_logs( + 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, + caplog, + ): + """Round-8 P1-4 at the real callback boundary: the row is emitted and + the payload-controlled exception message reaches neither the row nor + the application logs.""" + _ = mock_auth_default, mock_bq_client + import collections.abc + + class EvilMapping(collections.abc.Mapping): + + def __getitem__(self, k): + raise KeyError(k) + + def __len__(self): + return 1 + + def __iter__(self): + return iter(["a"]) + + def items(self): + raise RuntimeError("R8-CALLBACK-SECRET") + + 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={"hostile": EvilMapping()}, + ), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert "R8-CALLBACK-SECRET" not in json.dumps(log_entry, default=str) + assert "R8-CALLBACK-SECRET" not in caplog.text + attrs = json.loads(log_entry["attributes"]) + assert attrs["hostile"] == "[UNSUPPORTED_OBJECT]" + assert log_entry["is_truncated"] is True + + def test_scalar_key_collisions_fail_closed(self): + """Round-8 P2-10: scalar keys are normalized to their JSON form and + collisions get an explicit marker instead of silently collapsing.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + for pair in ( + {1: "n", "1": "s"}, + {True: "b", "true": "s"}, + {None: "x", "null": "s"}, + ): + out, truncated = truncate(pair, 10000) + assert truncated is True + assert len(out) == 2 + # Round-trip through JSON keeps both values. + assert len(json.loads(json.dumps(out))) == 2 + + # Round-9 P2-9: a pre-existing key in the marker namespace is never + # overwritten — markers are re-allocated until unique. + out, truncated = truncate( + {"[KEY_COLLISION_1]1": "reserved", "1": "string", 1: "numeric"}, + 10000, + ) + assert truncated is True + assert out["[KEY_COLLISION_1]1"] == "reserved" + assert out["1"] == "string" + assert len(out) == 3 + assert sorted(out.values()) == ["numeric", "reserved", "string"] + + def test_object_attr_traversal_bounded_and_selfref_terminates(self): + """Round-8 P2-8: __dict__ traversal charges the budget per entry and + self-references terminate immediately.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + class Big: + pass + + big = Big() + for i in range(200): + setattr(big, f"attr{i}", i) + out, truncated = truncate({"b": big}, 10000, None, 0, [50]) + assert truncated is True + assert len(out["b"]) <= 51 + + class Node: + pass + + node = Node() + node.self = node + node.access_token = "R8-SELF-SECRET" + out, _ = truncate({"n": node}, 10000) + assert out["n"]["self"] == "[CIRCULAR_REFERENCE]" + assert "R8-SELF-SECRET" not in json.dumps(out) + + def test_unlimited_mode_inspection_ceiling(self): + """Round-8 P2-13: max_content_length=-1 still bounds json.loads + materialization; over-ceiling container blobs fail closed.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + ceiling = bigquery_agent_analytics_plugin._MAX_JSON_INSPECT_CHARS + big = "[" + "1," * (ceiling // 2 + 10) + "1]" + out, truncated = truncate({"b": big}, -1) + assert out["b"] == "[UNPARSEABLE_JSON_BLOB]" + assert truncated is True + + @pytest.mark.asyncio + async def test_cancelled_shutdown_accounting_is_o1(self): + """Round-8 P2-12: external cancellation accounts queued rows without + a per-item synchronous drain.""" + + class CountingQueue(asyncio.Queue): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.get_nowait_calls = 0 + + def get_nowait(self): + self.get_nowait_calls += 1 + return super().get_nowait() + + bp = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="s", + batch_size=1, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=100, + shutdown_timeout=5.0, + ) + counting_queue = CountingQueue(maxsize=100) + bp._queue = counting_queue + + write_entered = asyncio.Event() + write_release = asyncio.Event() + + async def blocked_write(rows): + del rows + write_entered.set() + await write_release.wait() + + with mock.patch.object( + bp, "_write_rows_with_retry", side_effect=blocked_write + ): + await bp.start() + await bp.append({"r": 0}) + await write_entered.wait() + for i in range(3): + await bp.append({"r": i + 1}) + + closer = asyncio.create_task(bp.shutdown(timeout=30)) + await asyncio.sleep(0.05) + calls_before = counting_queue.get_nowait_calls + closer.cancel() + with pytest.raises(asyncio.CancelledError): + await closer + + # O(1): no per-item dequeue happened during cancellation — the queue + # was swapped out instead. + assert counting_queue.get_nowait_calls == calls_before + assert bp._queue is not counting_queue + assert bp.get_drop_stats()["shutdown_cancelled"] == 3 + + @pytest.mark.asyncio + async def test_aborted_setup_holds_rendezvous_and_allows_restart( + self, mock_auth_default, mock_bq_client + ): + """Round-8 P1-5: the setup rendezvous stays claimed until aborted + teardown completes, and a later restart fully succeeds.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + + entered = threading.Event() + release = threading.Event() + first_call = threading.Event() + + def gated_client(*args, **kwargs): + del args, kwargs + if not first_call.is_set(): + first_call.set() + entered.set() + release.wait(10) + return mock.MagicMock() + + future_held_during_teardown = [] + original_teardown = plugin._teardown_aborted_setup + + async def spying_teardown(): + future_held_during_teardown.append(plugin._setup_future is not None) + await original_teardown() + + write_client = mock.MagicMock() + write_client.transport = mock.MagicMock() + write_client.transport.close = mock.AsyncMock() + + with ( + mock.patch( + "google.adk.plugins.bigquery_agent_analytics_plugin.bigquery.Client", + side_effect=gated_client, + ), + mock.patch.object(plugin, "_teardown_aborted_setup", spying_teardown), + mock.patch.object( + bigquery_agent_analytics_plugin, + "BigQueryWriteAsyncClient", + return_value=write_client, + ), + mock.patch.object( + bigquery_agent_analytics_plugin.BatchProcessor, + "start", + mock.AsyncMock(), + ), + ): + owner = asyncio.create_task(plugin._ensure_started()) + while not entered.is_set(): + await asyncio.sleep(0.01) + await plugin.shutdown() + release.set() + outcome = await owner + assert outcome == "aborted" + # The rendezvous was still claimed while teardown ran, so no new + # setup could interleave and have its resources destroyed. + assert future_held_during_teardown == [True] + assert plugin._setup_future is None + + # A fresh start after the abort fully succeeds. + outcome2 = await plugin._ensure_started() + assert outcome2 == "ok" + assert plugin._started is True + assert plugin.client is not None + + @pytest.mark.asyncio + @pytest.mark.filterwarnings("error::RuntimeWarning") + @pytest.mark.filterwarnings("error::pytest.PytestUnraisableExceptionWarning") + async def test_host_timeout_effective_during_remote_drain( + self, mock_auth_default, mock_bq_client + ): + """Round-8 P1-6: a stuck remote loop must not block the event loop — + an outer host timeout fires instead of waiting out the full drain. + Warning-clean (#6360 round 9 P2-8): the shutdown coroutine created for + the never-running loop is explicitly closed, not leaked.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + + remote_loop = asyncio.new_event_loop() # never runs + try: + state = mock.MagicMock() + state.write_client = None + state.batch_processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="s", + batch_size=1, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=5.0, + ) + plugin._loop_state_by_loop[remote_loop] = state + + start = time.monotonic() + with pytest.raises(asyncio.TimeoutError): + await asyncio.wait_for(plugin.shutdown(timeout=5), timeout=0.1) + elapsed = time.monotonic() - start + # The old synchronous future.result(timeout=5) blocked the loop for + # the full remote timeout before the host timeout could fire. + assert elapsed < 2.0 + # The undrained state is retained for a retried close. + assert remote_loop in plugin._loop_state_by_loop + finally: + remote_loop.close() + + def test_drop_stats_stable_while_shutdown_folds( + self, mock_auth_default, mock_bq_client + ): + """Round-8 P2-7: claim+fold is one atomic transition, so readers never + observe a state as both live and folded (or neither).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + + async def scenario(): + gate = asyncio.Event() + + async def gated_processor_shutdown(timeout=None): + del timeout + await gate.wait() + + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock( + side_effect=gated_processor_shutdown + ) + state.batch_processor.get_drop_stats = mock.MagicMock( + return_value={"queue_full": 2} + ) + loop = asyncio.get_running_loop() + plugin._loop_state_by_loop[loop] = state + + closer = asyncio.create_task(plugin.shutdown(timeout=5)) + for _ in range(10): + await asyncio.sleep(0.005) + assert plugin.get_drop_stats().get("queue_full", 0) == 2 + gate.set() + await closer + assert plugin.get_drop_stats().get("queue_full", 0) == 2 + + asyncio.run(scenario()) + + @pytest.mark.asyncio + async def test_raced_event_counts_exactly_one_loss( + self, mock_auth_default, mock_bq_client, callback_context + ): + """Round-8 P2-9: one event racing shutdown records exactly one loss + (shutdown_race), not shutdown_race + setup_unavailable.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + with mock.patch.object( + plugin, "_ensure_started", mock.AsyncMock(return_value="aborted") + ): + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + stats = plugin.get_drop_stats() + assert stats.get("shutdown_race", 0) == 1 + assert stats.get("setup_unavailable", 0) == 0 + + @pytest.mark.asyncio + async def test_cancelled_setup_closes_eventual_client_and_executor( + self, mock_auth_default, mock_bq_client + ): + """Round-8 P2-14: cancelling a setup blocked in the client constructor + closes the eventual client and terminates the executor.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + + entered = threading.Event() + release = threading.Event() + eventual_client = mock.MagicMock() + + def gated_client(*args, **kwargs): + del args, kwargs + entered.set() + release.wait(10) + return eventual_client + + with mock.patch( + "google.adk.plugins.bigquery_agent_analytics_plugin.bigquery.Client", + side_effect=gated_client, + ): + owner = asyncio.create_task(plugin._ensure_started()) + while not entered.is_set(): + await asyncio.sleep(0.01) + executor = plugin._executor + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + release.set() + # The constructor thread finishes and the done-callback closes the + # orphaned client. + for _ in range(100): + if eventual_client.close.called: + break + await asyncio.sleep(0.02) + + assert eventual_client.close.called + assert plugin.client is None + assert plugin._executor is None + assert executor is not None and executor._shutdown + + @pytest.mark.asyncio + async def test_formatter_result_shapes_fail_closed( + 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, + caplog, + ): + """Round-9 P1-1: non-native formatter RESULTS fail closed inside the + boundary — a secret-returning or raising __str__ never reaches the + parser's str() fallback, the row, or the logs.""" + _ = mock_auth_default, mock_bq_client + + class LeakyResult: + + def __str__(self): + return "R9-FORMATTER-SECRET" + + class RaisingResult: + + def __str__(self): + raise RuntimeError("R9-FORMATTER-RAISE-SECRET") + + results = iter([LeakyResult(), RaisingResult()]) + + config = bigquery_agent_analytics_plugin.BigQueryLoggerConfig( + content_formatter=lambda content, event_type: next(results) + ) + 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): + for _ in range(2): + mock_write_client.append_rows.reset_mock() + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + dumped = json.dumps(log_entry, default=str) + assert "R9-FORMATTER-SECRET" not in dumped + assert "R9-FORMATTER-RAISE-SECRET" not in dumped + assert ( + bigquery_agent_analytics_plugin._FORMATTER_FAILED_SENTINEL + in log_entry["content"] + ) + assert "R9-FORMATTER-SECRET" not in caplog.text + assert "R9-FORMATTER-RAISE-SECRET" not in caplog.text + assert plugin.get_drop_stats().get("formatter_failed", 0) == 2 + + def test_value_backed_enums_not_published(self): + """Round-9 P1-2: StrEnum / (str, Enum) / bytes-backed members are + stringified through Enum.__str__ (member name), never their value.""" + import enum + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + class StrCred(str, enum.Enum): + access_token = "R9-STR-ENUM-SECRET" + + class BytesCred(bytes, enum.Enum): + token = b"R9-BYTES-ENUM-SECRET" + + payload = {"s": StrCred.access_token, "b": BytesCred.token} + if sys.version_info >= (3, 11): + + class NativeStrCred(enum.StrEnum): + refresh_token = "R9-STRENUM-SECRET" + + payload["n"] = NativeStrCred.refresh_token + + out, _ = truncate(payload, 10000) + dumped = json.dumps(out, default=str) + for canary in ( + "R9-STR-ENUM-SECRET", + "R9-BYTES-ENUM-SECRET", + "R9-STRENUM-SECRET", + ): + assert canary not in dumped, canary + assert out["s"] == "StrCred.access_token" + + def test_strip_bom_ws_is_linear(self): + """Round-9 P1-3: an alternating whitespace/BOM prefix is stripped in + one linear scan (the fixed-point slicing loop was quadratic).""" + strip = bigquery_agent_analytics_plugin._strip_bom_ws + prefix = " \ufeff" * 200_000 + start = time.monotonic() + assert strip(prefix + "{}") == "{}" + elapsed = time.monotonic() - start + # Quadratic behavior took minutes at this size; linear is ~25ms. + assert elapsed < 2.0 + + @pytest.mark.asyncio + async def test_failed_remote_drain_retains_state( + self, mock_auth_default, mock_bq_client, caplog + ): + """Round-9 P1-4: a remote drain that raises must NOT claim/fold its + state; it is retained for a retried close and the payload-controlled + message stays out of the logs.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + remote_loop = asyncio.new_event_loop() + thread = threading.Thread(target=remote_loop.run_forever, daemon=True) + thread.start() + try: + state = mock.MagicMock() + state.write_client = None + bp = mock.MagicMock(spec=bigquery_agent_analytics_plugin.BatchProcessor) + + async def failing_shutdown(timeout=None): + del timeout + raise RuntimeError("R9-DRAIN-SECRET") + + bp.shutdown = failing_shutdown + bp.get_drop_stats = mock.MagicMock(return_value={"queue_full": 1}) + state.batch_processor = bp + plugin._loop_state_by_loop[remote_loop] = state + + with caplog.at_level(logging.WARNING): + # Round-13 P1-3: a failed remote drain keeps teardown incomplete + # and surfaces to the owner instead of reporting success. + with pytest.raises( + bigquery_agent_analytics_plugin._ShutdownIncompleteError + ): + await plugin.shutdown(timeout=2) + + # Retained, not silently claimed as a successful drain. + assert remote_loop in plugin._loop_state_by_loop + assert plugin.get_drop_stats().get("queue_full", 0) == 1 + assert "R9-DRAIN-SECRET" not in caplog.text + assert "RuntimeError" in caplog.text + finally: + remote_loop.call_soon_threadsafe(remote_loop.stop) + thread.join(timeout=5) + remote_loop.close() + + @pytest.mark.asyncio + async def test_concurrent_shutdown_caller_awaits_completion( + self, mock_auth_default, mock_bq_client + ): + """Round-9 P1-5: a concurrent shutdown() caller coalesces on the + owner's completion instead of returning while teardown is running.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + gate = asyncio.Event() + + async def gated_shutdown(timeout=None): + del timeout + await gate.wait() + + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock(side_effect=gated_shutdown) + state.batch_processor.get_drop_stats = mock.MagicMock(return_value={}) + plugin._loop_state_by_loop[asyncio.get_running_loop()] = state + + first = asyncio.create_task(plugin.shutdown(timeout=5)) + await asyncio.sleep(0.05) + second = asyncio.create_task(plugin.shutdown()) + await asyncio.sleep(0.05) + assert not second.done(), "second caller returned mid-teardown" + gate.set() + await first + await asyncio.wait_for(second, timeout=5) + assert plugin._loop_state_by_loop == {} + + @pytest.mark.asyncio + async def test_shutdown_counts_rows_on_closed_loop( + self, mock_auth_default, mock_bq_client + ): + """Round-9 P1-6: queued rows owned by an already-closed loop are + counted as stale_loop when shutdown() claims the state.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + closed_loop = asyncio.new_event_loop() + closed_loop.close() + + bp = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="s", + batch_size=10, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=1.0, + ) + bp._queue.put_nowait({"r": 1}) + state = mock.MagicMock() + state.write_client = None + state.batch_processor = bp + plugin._loop_state_by_loop[closed_loop] = state + + await plugin.shutdown(timeout=1) + assert closed_loop not in plugin._loop_state_by_loop + assert plugin.get_drop_stats().get("stale_loop", 0) == 1 + + @pytest.mark.asyncio + async def test_completed_constructor_close_dispatched_off_loop( + self, mock_auth_default, mock_bq_client + ): + """Round-9 P1-7: when the constructor future is already done at + cancellation time, the orphan client's close still runs off-loop and + does not extend the cancellation window.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + + loop_thread_id = threading.get_ident() + close_started = threading.Event() + close_finished = threading.Event() + close_thread_ids = [] + eventual = mock.MagicMock() + + def slow_close(): + close_thread_ids.append(threading.get_ident()) + close_started.set() + time.sleep(0.2) + close_finished.set() + + eventual.close = slow_close + + def wrap_and_cancel(cf, **kwargs): + del kwargs + # Deterministic completed-before-cancel interleaving: wait for the + # constructor to finish, then deliver the cancellation. + cf.result(timeout=5) + raise asyncio.CancelledError() + + with ( + mock.patch( + "google.adk.plugins.bigquery_agent_analytics_plugin.bigquery.Client", + return_value=eventual, + ), + mock.patch.object( + bigquery_agent_analytics_plugin.asyncio, + "wrap_future", + side_effect=wrap_and_cancel, + ), + ): + start = time.monotonic() + with pytest.raises(asyncio.CancelledError): + await plugin._ensure_started() + elapsed = time.monotonic() - start + + assert close_started.wait(timeout=5) + assert close_finished.wait(timeout=5) + # The 200ms close did not run inline on the event-loop thread. + assert elapsed < 0.15 + assert close_thread_ids and close_thread_ids[0] != loop_thread_id + + def test_stale_cleanup_accounts_in_o1( + self, mock_auth_default, mock_bq_client + ): + """Round-9 P2-10: stale-loop cleanup accounts queued rows via qsize + minus sentinels, without a per-item synchronous drain.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + + class CountingQueue(asyncio.Queue): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.get_nowait_calls = 0 + + def get_nowait(self): + self.get_nowait_calls += 1 + return super().get_nowait() + + closed_loop = asyncio.new_event_loop() + closed_loop.close() + bp = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="s", + batch_size=10, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=2000, + shutdown_timeout=1.0, + ) + counting = CountingQueue(maxsize=2000) + for i in range(1000): + counting.put_nowait({"r": i}) + bp._queue = counting + state = mock.MagicMock() + state.write_client = None + state.batch_processor = bp + plugin._loop_state_by_loop[closed_loop] = state + + plugin._cleanup_stale_loop_states() + assert counting.get_nowait_calls == 0 + assert plugin.get_drop_stats().get("stale_loop", 0) == 1000 + + @pytest.mark.asyncio + async def test_shutdown_closes_shared_client( + self, mock_auth_default, mock_bq_client + ): + """Round-9 P2-11: normal shutdown closes the shared BigQuery client + instead of just dropping the reference.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + client = mock.MagicMock() + plugin.client = client + await plugin.shutdown(timeout=1) + assert client.close.called + assert plugin.client is None + + @pytest.mark.asyncio + async def test_cancelled_processor_shutdown_is_retryable(self): + """Round-7 P1-4: an externally cancelled BatchProcessor.shutdown() + (real processor, blocked writer) must not make later shutdown calls + re-raise the historical CancelledError; the retry completes and every + queued/in-flight row is accounted.""" + bp = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="s", + batch_size=1, + flush_interval=0.05, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=100, + shutdown_timeout=5.0, + ) + write_entered = asyncio.Event() + write_release = asyncio.Event() + + async def blocked_write(rows): + del rows + write_entered.set() + await write_release.wait() + + with mock.patch.object( + bp, "_write_rows_with_retry", side_effect=blocked_write + ): + await bp.start() + await bp.append({"r": 1}) + await write_entered.wait() # worker is blocked mid-write (in-flight=1) + for i in range(3): + await bp.append({"r": i + 2}) # three rows stay queued + + closer = asyncio.create_task(bp.shutdown(timeout=30)) + await asyncio.sleep(0.05) # let shutdown reach its wait_for + closer.cancel() + with pytest.raises(asyncio.CancelledError): + await closer + + # The retry completes instead of re-raising the historical + # cancellation, and nothing is left queued. + await bp.shutdown(timeout=1) + + assert bp._batch_processor_task.cancelled() + assert bp._queue.empty() + stats = bp.get_drop_stats() + # 3 queued rows (shutdown_cancelled) + 1 in-flight row counted by the + # cancelled worker (shutdown_timeout). + assert stats["shutdown_cancelled"] == 3 + assert stats["shutdown_timeout"] == 1 + + def test_overlimit_and_garbage_quoted_blobs_fail_closed(self): + """Round-7 P1-1: over-limit multi-layer quoted JSON and a valid quoted + credential layer with trailing garbage must not republish the secret.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + secret = "ROUND7-TRIPLE-SECRET-" + "x" * 300 + triple = json.dumps(json.dumps(json.dumps({"access_token": secret}))) + out, truncated = truncate({"blob": triple}, 64) + assert "ROUND7-TRIPLE-SECRET" not in json.dumps(out) + assert truncated + + trailing = ( + json.dumps(json.dumps({"access_token": "ROUND7-TRAIL-SECRET"})) + + " trailing" + ) + out, truncated = truncate({"blob": trailing}, 10000) + assert "ROUND7-TRAIL-SECRET" not in json.dumps(out) + # Round 8 P1-2 refined the policy: the leading string layer is + # redacted in place and the container-free suffix is preserved, so + # this is a redaction (changed), not a truncation. + assert "[REDACTED]" in out["blob"] + assert out["blob"].endswith(" trailing") + + # Ordinary quoted prose — with or without a suffix — stays untouched. + for prose in ('"hello" she said', '"just a quote"'): + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False + + @pytest.mark.asyncio + async def test_round7_shapes_redacted_at_row_boundary( + 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, + caplog, + ): + """Round-7 P1-2/P1-3/P2-8/P2-9 shapes at the final row boundary: the + row is always emitted, no canary reaches the row or the logs, + unsupported keys fail closed, and discarded binary reports + is_truncated.""" + _ = mock_auth_default, mock_bq_client + import collections + import types as types_module + + Cred = collections.namedtuple("Cred", ["access_token"]) + + class Trap: + + @property + def model_dump(self): + raise RuntimeError("ROUND7-PROPERTY-SECRET") + + class SneakyStr(str): + + def lstrip(self, *args): + return "plain" + + def startswith(self, *args): + return False + + 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={ + "named": Cred("ROUND7-NAMED-SECRET"), + "ns": types_module.SimpleNamespace( + access_token="ROUND7-REPR-SECRET", note="ok" + ), + "trap": Trap(), + "sneaky": SneakyStr('{"access_token":"ROUND7-STR-SUBCLASS"}'), + "bad_key": {(1, 2): "value"}, + "binary": b"\xff\xfe", + }, + ), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + dumped_row = json.dumps(log_entry, default=str) + for canary in ( + "ROUND7-NAMED-SECRET", + "ROUND7-REPR-SECRET", + "ROUND7-PROPERTY-SECRET", + "ROUND7-STR-SUBCLASS", + ): + assert canary not in dumped_row, canary + assert canary not in caplog.text, canary + attrs = json.loads(log_entry["attributes"]) + assert attrs["named"] == {"access_token": "[REDACTED]"} + assert attrs["ns"] == {"access_token": "[REDACTED]", "note": "ok"} + assert attrs["trap"] == "[UNSUPPORTED_OBJECT]" + assert attrs["bad_key"] == {"[UNSUPPORTED_KEY]": "value"} + assert attrs["binary"] == "[BINARY_DATA]" + assert log_entry["is_truncated"] is True + + 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() + outcome = await owner + + assert plugin._started is False + # Round-8 P2-9: the abort is reported structurally, not counted here; + # only a row owner converts it into a shutdown_race loss. + assert outcome == "aborted" + assert plugin.get_drop_stats().get("shutdown_race", 0) == 0 + + @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 + # And it routes through shutdown() semantics: counters remain queryable. + assert isinstance(plugin.get_drop_stats(), dict) + + @pytest.mark.asyncio + async def test_cancelled_close_releases_guard_and_allows_retry( + self, mock_auth_default, mock_bq_client + ): + """A close() cancelled mid-drain (PluginManager's close timeout) must + release _is_shutting_down, re-raise the cancellation, and leave the + retained loop state retryable by a second close (#6360 round 6 P1-1).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + loop = asyncio.get_running_loop() + + blocked = asyncio.Event() + release = asyncio.Event() + + async def blocking_shutdown(timeout=None): + del timeout + blocked.set() + await release.wait() + + state = mock.MagicMock() + state.write_client = None + state.batch_processor = mock.MagicMock( + spec=bigquery_agent_analytics_plugin.BatchProcessor + ) + state.batch_processor.shutdown = mock.AsyncMock( + side_effect=blocking_shutdown + ) + state.batch_processor.get_drop_stats = mock.MagicMock(return_value={}) + plugin._loop_state_by_loop[loop] = state + + closer = asyncio.create_task(plugin.close()) + await blocked.wait() + closer.cancel() + with pytest.raises(asyncio.CancelledError): + await closer + + # The guard is released and the undrained state is still retryable. + assert plugin._is_shutting_down is False + assert loop in plugin._loop_state_by_loop + + # A second close now completes and removes the retained state. + state.batch_processor.shutdown = mock.AsyncMock() + await plugin.close() + assert plugin._is_shutting_down is False + assert loop not in plugin._loop_state_by_loop + + @pytest.mark.asyncio + async def test_get_loop_state_uses_single_lookup( + self, mock_auth_default, mock_bq_client + ): + """A concurrent removal cannot split an existence check from lookup.""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + loop = asyncio.get_running_loop() + expected_state = mock.MagicMock() + + class DeleteOnContainsDict(dict): + + def __contains__(self, key): + present = super().__contains__(key) + if present: + del self[key] + return present + + plugin._loop_state_by_loop = DeleteOnContainsDict({loop: expected_state}) + + assert await plugin._get_loop_state() is expected_state + + @pytest.mark.asyncio + async def test_writer_built_during_shutdown_is_not_published( + self, mock_auth_default, mock_bq_client + ): + """A shutdown() that completes while _get_loop_state() is mid-build + must not let the fresh writer be published afterwards; the new + processor and transport are torn down instead (#6360 round 6 P1-2).""" + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + + start_entered = asyncio.Event() + start_gate = asyncio.Event() + processor_shutdowns = [] + + async def gated_start(self): + del self + start_entered.set() + await start_gate.wait() + + async def record_shutdown(self, timeout=None): + del self + processor_shutdowns.append(timeout) + + transport = mock.MagicMock() + transport.close = mock.AsyncMock() + write_client = mock.MagicMock() + write_client.transport = transport + + with ( + mock.patch.object( + bigquery_agent_analytics_plugin.BatchProcessor, + "start", + gated_start, + ), + mock.patch.object( + bigquery_agent_analytics_plugin.BatchProcessor, + "shutdown", + record_shutdown, + ), + mock.patch.object( + bigquery_agent_analytics_plugin, + "BigQueryWriteAsyncClient", + return_value=write_client, + ), + ): + builder = asyncio.create_task(plugin._get_loop_state()) + await start_entered.wait() + # shutdown() completes while the writer is still being built: its + # snapshot is empty, so only the publication guard can stop the leak. + await plugin.shutdown() + start_gate.set() + with pytest.raises(RuntimeError): + await builder + + assert plugin._loop_state_by_loop == {} + assert processor_shutdowns, "fresh processor must be shut down" + transport.close.assert_awaited() + + 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_double_encoded_and_rootmodel_blobs_are_redacted(self): + """Round-6 P1-3 shapes: JSON-encoded string layers (double/triple + json.dumps) and scalar model_dump() results (RootModel[str]) re-enter + the redaction path instead of bypassing it.""" + import pydantic + + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + double = json.dumps(json.dumps({"access_token": "DOUBLE-ENCODED-SECRET"})) + out, _ = truncate({"blob": double}, 10000) + assert "DOUBLE-ENCODED-SECRET" not in json.dumps(out) + + triple = json.dumps(double) + out, _ = truncate({"blob": triple}, 10000) + assert "DOUBLE-ENCODED-SECRET" not in json.dumps(out) + + root = pydantic.RootModel[str]('{"access_token":"ROOT-SECRET"}') + out, _ = truncate({"model": root}, 10000) + assert "ROOT-SECRET" not in json.dumps(out) + + # Ordinary quoted prose (not a JSON document) is left untouched. + prose = '"hello" she said' + out, truncated = truncate({"s": prose}, 10000) + assert out["s"] == prose + assert truncated is False + + def test_depth_truncated_json_blob_reports_truncation(self): + """Round-6 P2-6: a JSON blob rewritten with [MAX_DEPTH_EXCEEDED] + discards payload and must therefore report truncated=True.""" + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + deep = "[" * 60 + "]" * 60 + out, truncated = truncate({"blob": deep}, 10000) + assert "[MAX_DEPTH_EXCEEDED]" in out["blob"] + assert truncated is True + + def test_sanitizer_stops_at_node_budget(self): + """A very wide payload stops at the work budget, emits ONE remainder + sentinel, and the output stays bounded by the budget — iteration used + to continue over the full input, appending one sentinel per remaining + element (#6360 review round 5 P2, round 6 P2-5).""" + max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + wide = list(range(max_nodes * 2)) + out, truncated = truncate({"wide": wide}, 10000) + assert truncated + assert out["wide"][-1] == "[SANITIZE_BUDGET_EXCEEDED]" + assert out["wide"].count("[SANITIZE_BUDGET_EXCEEDED]") == 1 + # Bounded output: budget entries plus the single remainder sentinel. + assert len(out["wide"]) <= max_nodes + 1 + + def test_sanitizer_budget_covers_directly_redacted_entries(self): + """Directly redacted keys (temp:/sensitive) consume budget too — a + wide temp: mapping used to bypass the bound entirely and report + truncated=False (#6360 review round 6 P2-5).""" + max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES + truncate = bigquery_agent_analytics_plugin._recursive_smart_truncate + + wide_temp = {f"temp:{i}": i for i in range(max_nodes * 2)} + out, truncated = truncate(wide_temp, 10000) + assert truncated + assert len(out) <= max_nodes + 1 + assert "[SANITIZE_BUDGET_EXCEEDED]" in out + + @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 + + +class TestLatestReviewLifecycleRegressions: + """Regressions for lifecycle findings in PR #11 review 4731058331.""" + + @pytest.mark.asyncio + async def test_later_before_model_short_circuit_does_not_leak_span( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + invocation_context, + dummy_arrow_schema, + ): + """A synthesized response row belongs to the parent agent span.""" + + class ShortCircuitPlugin(bigquery_agent_analytics_plugin.BasePlugin): + + def __init__(self): + super().__init__(name="short_circuit") + + async def before_model_callback(self, **kwargs): + del kwargs + return llm_response_lib.LlmResponse( + content=types.Content( + role="model", parts=[types.Part(text="cached response")] + ) + ) + + trace_manager = bigquery_agent_analytics_plugin.TraceManager + trace_manager.clear_stack() + try: + parent_span_id = trace_manager.push_span(callback_context, "agent") + manager = plugin_manager_lib.PluginManager( + [bq_plugin_inst, ShortCircuitPlugin()] + ) + request = llm_request_lib.LlmRequest( + model="gemini-pro", + contents=[ + types.Content(role="user", parts=[types.Part(text="prompt")]) + ], + ) + + short_response = await manager.run_before_model_callback( + callback_context=callback_context, llm_request=request + ) + assert short_response is not None + leaked_span_id = trace_manager.get_current_span_id() + assert leaked_span_id != parent_span_id + + # ADK intentionally skips run_after_model_callback on this path and + # emits the synthesized response as a non-partial event instead. + await bq_plugin_inst.flush() + mock_write_client.append_rows.reset_mock() + event = event_lib.Event( + author="agent", + content=short_response.content, + ) + await manager.run_on_event_callback( + invocation_context=invocation_context, event=event + ) + await bq_plugin_inst.flush() + + rows = await _get_captured_rows_async( + mock_write_client, dummy_arrow_schema + ) + response_row = next( + row for row in rows if row["event_type"] == "AGENT_RESPONSE" + ) + assert response_row["span_id"] == parent_span_id + assert response_row["span_id"] != leaked_span_id + assert trace_manager.get_current_span_id() == parent_span_id + finally: + trace_manager.clear_stack() + + @pytest.mark.asyncio + async def test_partial_event_preserves_live_llm_span( + self, bq_plugin_inst, callback_context, invocation_context + ): + trace_manager = bigquery_agent_analytics_plugin.TraceManager + trace_manager.clear_stack() + try: + trace_manager.push_span(callback_context, "agent") + await bq_plugin_inst.before_model_callback( + callback_context=callback_context, + llm_request=llm_request_lib.LlmRequest(model="gemini-pro"), + ) + llm_span_id = trace_manager.get_current_span_id() + + await bq_plugin_inst.on_event_callback( + invocation_context=invocation_context, + event=event_lib.Event( + author="agent", + partial=True, + content=types.Content( + role="model", parts=[types.Part(text="stream chunk")] + ), + ), + ) + assert trace_manager.get_current_span_id() == llm_span_id + finally: + trace_manager.clear_stack() + + @pytest.mark.asyncio + async def test_external_cancel_during_worker_ack_is_not_swallowed(self): + processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="stream", + batch_size=1, + flush_interval=1.0, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=1.0, + ) + first_cancel_seen = asyncio.Event() + never = asyncio.Event() + + async def slow_cancel_ack(): + try: + await never.wait() + except asyncio.CancelledError: + first_cancel_seen.set() + await never.wait() + + processor._batch_processor_task = asyncio.create_task(slow_cancel_ack()) + processor._queue.put_nowait({"row": 1}) + + with mock.patch.object( + bigquery_agent_analytics_plugin.asyncio, + "wait_for", + new=mock.AsyncMock(side_effect=asyncio.TimeoutError), + ): + owner = asyncio.create_task(processor.shutdown(timeout=0.01)) + await first_cancel_seen.wait() + owner.cancel() + with pytest.raises(asyncio.CancelledError): + await owner + + # A later close owns the retained terminal task/queue and accounts it. + await processor.shutdown(timeout=1.0) + assert processor._batch_processor_task.cancelled() + assert processor._queue.empty() + assert processor.get_drop_stats()["shutdown_timeout"] == 1 + + @pytest.mark.asyncio + async def test_dead_loop_state_is_replaced_once_and_rows_are_accounted( + self, mock_auth_default, mock_bq_client + ): + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + plugin._write_stream_name = DEFAULT_STREAM_NAME + + old_transport = mock.MagicMock() + close_entered = asyncio.Event() + close_release = asyncio.Event() + + async def gated_old_close(): + close_entered.set() + await close_release.wait() + + old_transport.close = mock.AsyncMock(side_effect=gated_old_close) + old_client = mock.MagicMock(transport=old_transport) + old_processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=old_client, + arrow_schema=None, + write_stream=DEFAULT_STREAM_NAME, + batch_size=1, + flush_interval=0.01, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=1.0, + ) + old_processor._shutdown = True + old_processor._dropped["queue_full"] = 2 + old_processor._queue.put_nowait({"old": 1}) + old_processor._queue.put_nowait({"old": 2}) + old_processor._batch_processor_task = asyncio.create_task(asyncio.sleep(0)) + await old_processor._batch_processor_task + + loop = asyncio.get_running_loop() + old_state = bigquery_agent_analytics_plugin._LoopState( + old_client, old_processor + ) + plugin._loop_state_by_loop[loop] = old_state + + new_transport = mock.MagicMock() + new_transport.close = mock.AsyncMock() + new_client = mock.MagicMock(transport=new_transport) + with mock.patch.object( + bigquery_agent_analytics_plugin, + "BigQueryWriteAsyncClient", + return_value=new_client, + ): + builder = asyncio.create_task(plugin._get_loop_state()) + await close_entered.wait() + + # Replacement is already published while the sole owner closes the old + # transport, so a concurrent caller cannot build a second processor. + concurrent_state = await plugin._get_loop_state() + close_release.set() + replacement = await builder + + assert replacement is concurrent_state + assert replacement is plugin._loop_state_by_loop[loop] + assert replacement is not old_state + old_transport.close.assert_awaited_once() + stats = plugin.get_drop_stats() + assert stats["queue_full"] == 2 + assert stats["shutdown_timeout"] == 2 + + write_rows = mock.AsyncMock() + with mock.patch.object( + replacement.batch_processor, + "_write_rows_with_retry", + new=write_rows, + ): + row = {"new": 1} + await replacement.batch_processor.append(row) + await asyncio.wait_for(replacement.batch_processor.flush(), timeout=1) + write_rows.assert_awaited_once_with([row]) + + await plugin.shutdown(timeout=1) + + @pytest.mark.asyncio + async def test_normal_timeout_retrieves_worker_cancel_and_drains_queue(self): + """The 3.10-compatible acknowledgement path handles worker cancel.""" + processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="stream", + batch_size=1, + flush_interval=1.0, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=1.0, + ) + never = asyncio.Event() + processor._batch_processor_task = asyncio.create_task(never.wait()) + processor._queue.put_nowait({"row": 1}) + + await processor.shutdown(timeout=0.001) + + assert processor._batch_processor_task.cancelled() + assert processor._queue.empty() + assert processor.get_drop_stats()["shutdown_timeout"] == 1 + + @pytest.mark.asyncio + async def test_error_columns_and_tracebacks_redact_embedded_credentials( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + secrets = ( + "AUTH-SECRET", + "QUERY-SECRET", + "JSON-SECRET", + "SIGNATURE-SECRET", + "ESCAPED-SECRET", + ) + message = ( + "safe prefix Authorization: Bearer AUTH-SECRET; " + "access-token=QUERY-SECRET" + ) + traceback_text = ( + 'Traceback safe prefix {"access_token":"JSON-SECRET"}; ' + 'next {"access\\u005ftoken":"ESCAPED-SECRET"}; ' + "x-goog-signature=SIGNATURE-SECRET" + ) + + await bq_plugin_inst._log_event( + "AGENT_ERROR", + callback_context, + raw_content={"error_traceback": traceback_text}, + event_data=bigquery_agent_analytics_plugin.EventData( + status="ERROR", error_message=message + ), + ) + await bq_plugin_inst.flush() + row = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + stored = json.dumps(row, default=str) + assert all(secret not in stored for secret in secrets) + assert row["error_message"].startswith("safe prefix Authorization:") + assert bigquery_agent_analytics_plugin._sanitize_sensitive_text( + "Authorization: Bearer AUTH-SECRET", -1 + ) == ("Authorization: [REDACTED]", True) + for escaped in ( + r"access\u005ftoken=ESCAPED-SECRET", + r'Traceback { "access\u005ftoken":"ESCAPED-SECRET"}', + ): + assert bigquery_agent_analytics_plugin._sanitize_sensitive_text( + escaped, -1 + ) == ("[REDACTED_SENSITIVE_TEXT]", True) + assert bigquery_agent_analytics_plugin._sanitize_sensitive_text( + "temp:credential=TEMP-SECRET", -1 + ) == ("temp:credential=[REDACTED]", True) + assert "[REDACTED]" in stored + assert row["is_truncated"] is True + + @pytest.mark.asyncio + async def test_safe_error_message_is_preserved_exactly( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + message = "[INFO] ordinary failure at worker 7" + await bq_plugin_inst._log_event( + "LLM_ERROR", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData( + status="ERROR", error_message=message + ), + ) + await bq_plugin_inst.flush() + row = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + assert row["error_message"] == message + assert row["is_truncated"] is False + + def test_sensitive_text_redacts_complete_values_and_encoded_constructs(self): + sanitize = bigquery_agent_analytics_plugin._sanitize_sensitive_text + redacted_cases = { + "access_token=[REDACTED]SECRET": "SECRET", + "access_token=[REDACTED]]SECRET": "SECRET", + "access_token=[REDACTED]/SECRET": "SECRET", + 'Authorization: Digest username="u", response="DIGEST-SECRET"': ( + "DIGEST-SECRET" + ), + ( + "Authorization: AWS4-HMAC-SHA256 " + "Credential=AWS-SECRET, SignedHeaders=host" + ): "AWS-SECRET", + "Proxy-Authorization: Negotiate NEGOTIATE-SECRET": "NEGOTIATE-SECRET", + "Bearer\nBEARER-SECRET": "BEARER-SECRET", + "Basic\tdXNlcjpwYXNz": "dXNlcjpwYXNz", + "Basic dXNlcg==": "dXNlcg==", + "sig=SIG-SECRET": "SIG-SECRET", + "x-amz-signature=AMZ-SIGNATURE-SECRET": "AMZ-SIGNATURE-SECRET", + "x_amz_credential=AMZ-CREDENTIAL-SECRET": "AMZ-CREDENTIAL-SECRET", + "google-access-id=GOOGLE-ID-SECRET": "GOOGLE-ID-SECRET", + r"access\u005ftoken=UNICODE-SECRET": "UNICODE-SECRET", + r"access\x5ftoken=HEX-SECRET": "HEX-SECRET", + "access_token%3DPERCENT-SECRET": "PERCENT-SECRET", + "access%255Ftoken%253DDOUBLE-SECRET": "DOUBLE-SECRET", + } + for value, secret in redacted_cases.items(): + sanitized, changed = sanitize(value, -1) + assert changed is True, value + assert secret not in sanitized, value + + # A sentinel is idempotent only when it is the complete value. + assert sanitize("access_token=[REDACTED]", -1) == ( + "access_token=[REDACTED]", + False, + ) + + structured, _ = bigquery_agent_analytics_plugin._recursive_smart_truncate( + { + "x-amz-signature": "STRUCTURED-AMZ-SECRET", + "google_access_id": "STRUCTURED-GOOGLE-SECRET", + "safe": True, + }, + -1, + ) + assert structured == { + "x-amz-signature": "[REDACTED]", + "google_access_id": "[REDACTED]", + "safe": True, + } + + def test_sensitive_text_preserves_safe_slashes_and_encoded_prose_exactly( + self, + ): + sanitize = bigquery_agent_analytics_plugin._sanitize_sensitive_text + safe = ( + r"C:\Users\secret\project\file.json", + r"Invalid \escape at position 4", + r"can't decode \x5c in position 2", + "the bearer of bad news", + "a basic principle", + "a basic test", + "design=balanced", + "signal=strong", + "progress%3D100%25 complete", + "literal%2525value", + ) + for value in safe: + assert sanitize(value, -1) == (value, False) + + # A moderately wide safe input exercises the bounded stack scanner while + # pinning the useful property instead of a timing threshold. + wide = (r"C:\safe\secret\file%25.txt; " * 20_000).rstrip() + assert sanitize(wide, len(wide)) == (wide, False) + + @pytest.mark.asyncio + async def test_live_shutting_down_writer_aborts_owner_and_waiter_once( + self, mock_auth_default, mock_bq_client, callback_context + ): + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(), + arrow_schema=None, + write_stream="stream", + batch_size=1, + flush_interval=1.0, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=1.0, + ) + never = asyncio.Event() + processor._batch_processor_task = asyncio.create_task(never.wait()) + processor._shutdown = True + loop = asyncio.get_running_loop() + plugin._loop_state_by_loop[loop] = ( + bigquery_agent_analytics_plugin._LoopState(mock.MagicMock(), processor) + ) + + entered = asyncio.Event() + release = asyncio.Event() + + async def attempt_setup(**kwargs): + del kwargs + entered.set() + await release.wait() + await plugin._get_loop_state() + + try: + with mock.patch.object(plugin, "_lazy_setup", side_effect=attempt_setup): + owner = asyncio.create_task(plugin._ensure_started()) + await entered.wait() + waiter = asyncio.create_task(plugin._ensure_started()) + await asyncio.sleep(0) + release.set() + assert await owner == "aborted" + assert await waiter == "aborted" + + assert plugin._startup_error is None + assert plugin._setup_failures == 0 + assert plugin._setup_retry_at == 0 + assert plugin._loop_state_by_loop[loop].batch_processor is processor + + await plugin._log_event( + "STATE_DELTA", + callback_context, + event_data=bigquery_agent_analytics_plugin.EventData(), + ) + assert plugin.get_drop_stats()["shutdown_race"] == 1 + finally: + processor._batch_processor_task.cancel() + with contextlib.suppress(asyncio.CancelledError): + await processor._batch_processor_task + + @pytest.mark.asyncio + async def test_detached_transport_closed_when_replacement_build_fails( + self, mock_auth_default, mock_bq_client + ): + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + loop = asyncio.get_running_loop() + + old_transport = mock.MagicMock(close=mock.AsyncMock()) + old_processor = bigquery_agent_analytics_plugin.BatchProcessor( + write_client=mock.MagicMock(transport=old_transport), + arrow_schema=None, + write_stream="stream", + batch_size=1, + flush_interval=1.0, + retry_config=bigquery_agent_analytics_plugin.RetryConfig(), + queue_max_size=10, + shutdown_timeout=1.0, + ) + old_processor._shutdown = True + old_processor._batch_processor_task = asyncio.create_task(asyncio.sleep(0)) + await old_processor._batch_processor_task + plugin._loop_state_by_loop[loop] = ( + bigquery_agent_analytics_plugin._LoopState( + mock.MagicMock(transport=old_transport), old_processor + ) + ) + + fresh_transport = mock.MagicMock(close=mock.AsyncMock()) + fresh_client = mock.MagicMock(transport=fresh_transport) + with ( + mock.patch.object( + bigquery_agent_analytics_plugin, + "BigQueryWriteAsyncClient", + return_value=fresh_client, + ), + mock.patch.object( + bigquery_agent_analytics_plugin.BatchProcessor, + "__init__", + side_effect=RuntimeError("construction failed"), + ), + ): + with pytest.raises(RuntimeError, match="construction failed"): + await plugin._get_loop_state() + + old_transport.close.assert_awaited_once() + fresh_transport.close.assert_awaited_once() + assert loop not in plugin._loop_state_by_loop + + @pytest.mark.asyncio + async def test_invalidated_writer_transport_closes_when_shutdown_is_cancelled( + self, mock_auth_default, mock_bq_client + ): + _ = mock_auth_default, mock_bq_client + plugin = bigquery_agent_analytics_plugin.BigQueryAgentAnalyticsPlugin( + PROJECT_ID, DATASET_ID, table_id=TABLE_ID + ) + plugin._credentials = mock.MagicMock(quota_project_id=None) + + start_entered = asyncio.Event() + start_release = asyncio.Event() + shutdown_entered = asyncio.Event() + shutdown_never = asyncio.Event() + + async def gated_start(self): + del self + start_entered.set() + await start_release.wait() + + async def blocked_shutdown(self, timeout=None): + del self, timeout + shutdown_entered.set() + await shutdown_never.wait() + + transport = mock.MagicMock(close=mock.AsyncMock()) + write_client = mock.MagicMock(transport=transport) + with ( + mock.patch.object( + bigquery_agent_analytics_plugin.BatchProcessor, + "start", + gated_start, + ), + mock.patch.object( + bigquery_agent_analytics_plugin.BatchProcessor, + "shutdown", + blocked_shutdown, + ), + mock.patch.object( + bigquery_agent_analytics_plugin, + "BigQueryWriteAsyncClient", + return_value=write_client, + ), + ): + builder = asyncio.create_task(plugin._get_loop_state()) + await start_entered.wait() + await plugin.shutdown() + start_release.set() + await shutdown_entered.wait() + builder.cancel() + with pytest.raises(asyncio.CancelledError): + await builder + + transport.close.assert_awaited_once() + assert plugin._loop_state_by_loop == {} + + @pytest.mark.asyncio + async def test_raw_bracket_prose_preserved_inline_and_gcs(self): + inline = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + prose = ("[INFO] ready", "[link](https://example.test)", "{not json}") + for value in prose: + payload, parts, truncated = await inline.parse( + types.Content(parts=[types.Part(text=value)]) + ) + assert payload == {"text_summary": value} + assert parts[0]["text"] == value + assert truncated is False + + offloader = mock.AsyncMock() + offloader.upload_content.return_value = "gs://bucket/safe.txt" + offloaded = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=offloader, + trace_id="t", + span_id="s", + max_length=-1, + ) + large_prose = "[INFO] " + "safe prose " * 4000 + await offloaded.parse(types.Content(parts=[types.Part(text=large_prose)])) + assert offloader.upload_content.call_args.args[0] == large_prose + + @pytest.mark.asyncio + async def test_bracket_prose_auth_and_signature_classification(self): + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + safe = ( + r"[INFO] C:\Users\secret\project", + "[INFO] the bearer of bad news", + "[INFO] a basic principle", + "[INFO] a basic test", + "[INFO] design=balanced and progress%3D100%25", + ) + for value in safe: + payload, parts, truncated = await parser.parse( + types.Content(parts=[types.Part(text=value)]) + ) + assert payload == {"text_summary": value} + assert parts[0]["text"] == value + assert truncated is False + + unsafe = ( + ("[WARN] Bearer\tBRACKET-BEARER-SECRET", "BRACKET-BEARER-SECRET"), + ("[WARN] Basic\ndXNlcjpwYXNz", "dXNlcjpwYXNz"), + ("[WARN] sig=BRACKET-SIG-SECRET", "BRACKET-SIG-SECRET"), + ( + "[WARN] x-amz-signature=BRACKET-AMZ-SECRET", + "BRACKET-AMZ-SECRET", + ), + ( + "[WARN] access%255Ftoken%253DBRACKET-ENCODED-SECRET", + "BRACKET-ENCODED-SECRET", + ), + ) + for value, secret in unsafe: + payload, parts, truncated = await parser.parse( + types.Content(parts=[types.Part(text=value)]) + ) + stored = json.dumps({"payload": payload, "parts": parts}) + assert secret not in stored + assert "[UNPARSEABLE_JSON_BLOB]" in stored + assert truncated is True + + @pytest.mark.asyncio + async def test_malformed_bracket_credentials_still_fail_closed(self): + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + for value in ( + '{"access_token":"MALFORMED-SECRET"', + '{"access\\u005ftoken":"ESCAPED-SECRET"', + ): + payload, parts, truncated = await parser.parse( + types.Content(parts=[types.Part(text=value)]) + ) + stored = json.dumps({"payload": payload, "parts": parts}) + assert "MALFORMED-SECRET" not in stored + assert "ESCAPED-SECRET" not in stored + assert "[UNPARSEABLE_JSON_BLOB]" in stored + assert truncated is True + + @pytest.mark.asyncio + async def test_external_uri_redacts_query_fragment_and_userinfo(self): + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + signed = ( + "https://storage.example.test/safe/path?safe=kept" + "&X-Goog-Credential=URI-CREDENTIAL" + "&X-Goog-Signature=URI-SIGNATURE#access-token=FRAGMENT-SECRET" + ) + _, parts, truncated = await parser.parse( + types.Content( + parts=[types.Part.from_uri(file_uri=signed, mime_type="text/plain")] + ) + ) + uri = parts[0]["uri"] + assert uri.startswith("https://storage.example.test/safe/path?") + assert "safe=kept" in uri + assert all( + secret not in uri + for secret in ("URI-CREDENTIAL", "URI-SIGNATURE", "FRAGMENT-SECRET") + ) + assert truncated is True + + userinfo = types.Part( + file_data=types.FileData( + file_uri="https://user:password@example.test/safe", + mime_type="text/plain", + ) + ) + _, parts, truncated = await parser.parse(types.Content(parts=[userinfo])) + assert parts[0]["uri"] == "[REDACTED_SENSITIVE_URI]" + assert truncated is True + + @pytest.mark.asyncio + async def test_external_uri_redacts_sensitive_path_segments_and_variants( + self, + ): + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + uri = ( + "https://example.test/public/access-token/PATH-SECRET/report" + "?x-amz-signature=QUERY-SIGNATURE-SECRET" + "&access%255Ftoken%253DDOUBLE-QUERY-SECRET" + ) + _, parts, truncated = await parser.parse( + types.Content( + parts=[types.Part.from_uri(file_uri=uri, mime_type="text/plain")] + ) + ) + stored_uri = parts[0]["uri"] + for secret in ( + "PATH-SECRET", + "QUERY-SIGNATURE-SECRET", + "DOUBLE-QUERY-SECRET", + ): + assert secret not in stored_uri + assert "/public/%5BREDACTED%5D/%5BREDACTED%5D/report" in stored_uri + assert truncated is True + + safe_uri = "https://example.test/design/signal/public/progress%25/report" + _, parts, truncated = await parser.parse( + types.Content( + parts=[ + types.Part.from_uri(file_uri=safe_uri, mime_type="text/plain") + ] + ) + ) + assert parts[0]["uri"] == safe_uri + assert truncated is False + + missing = types.Part( + file_data=types.FileData(file_uri=None, mime_type="text/plain") + ) + _, parts, truncated = await parser.parse(types.Content(parts=[missing])) + assert parts[0]["uri"] == "[REDACTED_SENSITIVE_URI]" + assert truncated is True + + @pytest.mark.asyncio + async def test_structured_non_text_parts_are_complete_and_private(self): + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=1000, + ) + secret = "STRUCTURED-PART-SECRET" + content = types.Content( + parts=[ + types.Part( + function_response=types.FunctionResponse( + name="lookup", response={"access_token": secret, "ok": True} + ) + ), + types.Part( + executable_code=types.ExecutableCode( + language=types.Language.PYTHON, + code=json.dumps({"private_key": secret}), + ) + ), + types.Part( + code_execution_result=types.CodeExecutionResult( + outcome=types.Outcome.OUTCOME_OK, + output=f"Authorization: Bearer {secret}", + ) + ), + ] + ) + + payload, parts, truncated = await parser.parse(content) + stored = json.dumps({"payload": payload, "parts": parts}) + assert secret not in stored + assert truncated is True + assert "Function response: lookup" in payload["text_summary"] + assert "Executable code" in payload["text_summary"] + assert "Code execution result" in payload["text_summary"] + assert "function_response" in json.loads(parts[0]["part_attributes"]) + assert "executable_code" in json.loads(parts[1]["part_attributes"]) + assert "code_execution_result" in json.loads(parts[2]["part_attributes"]) + + def test_structured_part_dictionary_keys_are_sanitized_without_collisions( + self, + ): + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=1000, + ) + value = mock.MagicMock() + value.model_dump.return_value = { + "access_token=[REDACTED]": "genuine-marker", + "access_token=KEY-ONE-SECRET": "first", + "access-token=KEY-TWO-SECRET": "second", + "[KEY_COLLISION_1]access_token=[REDACTED]": "genuine-collision", + "sig": "STRUCTURED-SIG-SECRET", + "x-amz-credential": "STRUCTURED-AMZ-SECRET", + } + + serialized, content_lost = parser._serialize_part_model(value) + dumped = json.dumps(serialized) + assert "KEY-ONE-SECRET" not in dumped + assert "KEY-TWO-SECRET" not in dumped + assert "STRUCTURED-SIG-SECRET" not in dumped + assert "STRUCTURED-AMZ-SECRET" not in dumped + assert sorted(serialized.values()) == [ + "[REDACTED]", + "[REDACTED]", + "first", + "genuine-collision", + "genuine-marker", + "second", + ] + assert len(serialized) == 6 + assert any(key.startswith("[KEY_COLLISION_") for key in serialized) + assert content_lost is True