diff --git a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py index 3a09fc942f8..97eb5845ab2 100644 --- a/src/google/adk/plugins/bigquery_agent_analytics_plugin.py +++ b/src/google/adk/plugins/bigquery_agent_analytics_plugin.py @@ -43,6 +43,8 @@ from typing import Callable from typing import Optional from typing import TYPE_CHECKING +from urllib.parse import urlsplit +from urllib.parse import urlunsplit import uuid import weakref @@ -268,9 +270,33 @@ def _get_tool_origin( "password", }) +# Written to the content column in place of the payload when the configured +# content_formatter raises. +_FORMATTER_FAILED_SENTINEL = "[FORMATTER_FAILED]" + +# Written in place of an external URI that cannot be stored safely. +_REDACTED_URI = "[REDACTED_SENSITIVE_URI]" + +# A URI longer than this is replaced wholesale rather than parsed. +_MAX_URI_LENGTH = 8192 + +# Deepest nesting _recursive_smart_truncate walks before replacing a value. +_MAX_SANITIZE_DEPTH = 50 + +# Total nodes one sanitizer invocation may visit. Depth and per-string size +# are bounded, but width was not: a million-element list, or an object that +# manufactures two fresh children per access, walks tens of millions of +# nodes inside the depth cap. The remainder is replaced with a sentinel and +# the row is flagged truncated. +_MAX_SANITIZE_NODES = 100_000 + def _recursive_smart_truncate( - obj: Any, max_len: int, seen: Optional[set[int]] = None + obj: Any, + max_len: int, + seen: Optional[set[int]] = None, + depth: int = 0, + budget: Optional[list[int]] = None, ) -> tuple[Any, bool]: """Recursively truncates string values within a dict or list. @@ -281,12 +307,29 @@ def _recursive_smart_truncate( obj: The object to truncate. max_len: Maximum length for string values. seen: Set of object IDs visited in the current recursion stack. + depth: Current recursion depth. + budget: Single-element list holding the nodes left in the shared work + budget for this invocation. Returns: A tuple of (truncated_object, is_truncated). """ if seen is None: seen = set() + if budget is None: + budget = [_MAX_SANITIZE_NODES] + budget[0] -= 1 + if budget[0] < 0: + return "[SANITIZE_BUDGET_EXCEEDED]", True + + # The id()-based cycle detection below cannot catch an object graph that + # manufactures a new object on every duck-typed access, which is what any + # object whose model_dump()/dict()/to_dict() returns a fresh wrapper does. + # Such a graph recurses until the interpreter's own limit. The replacement + # discards real data, so unlike "[CIRCULAR_REFERENCE]" it reports + # truncation. + if depth >= _MAX_SANITIZE_DEPTH: + return "[MAX_DEPTH_EXCEEDED]", True obj_id = id(obj) if obj_id in seen: @@ -315,13 +358,26 @@ def _recursive_smart_truncate( # but explicit loop is fine for clarity given recursive nature. new_dict = {} for k, v in obj.items(): + # Stop iterating once the budget is exhausted. Recursing on every + # remaining entry still did work proportional to the input and + # produced one sentinel per entry; a single remainder sentinel + # stands in for everything dropped. + if budget[0] <= 0: + new_dict["[SANITIZE_BUDGET_EXCEEDED]"] = "[SANITIZE_BUDGET_EXCEEDED]" + truncated_any = True + break if isinstance(k, str): k_lower = k.lower() if k_lower in _SENSITIVE_KEYS or k_lower.startswith("temp:"): + # A directly redacted entry costs budget too, otherwise a wide + # "temp:" mapping bypasses the bound entirely. + budget[0] -= 1 new_dict[k] = "[REDACTED]" continue - val, trunc = _recursive_smart_truncate(v, max_len, seen) + val, trunc = _recursive_smart_truncate( + v, max_len, seen, depth + 1, budget + ) if trunc: truncated_any = True new_dict[k] = val @@ -331,7 +387,14 @@ def _recursive_smart_truncate( new_list = [] # Explicit loop to handle flag propagation for i in obj: - val, trunc = _recursive_smart_truncate(i, max_len, seen) + # Same bound as the mapping loop. + if budget[0] <= 0: + new_list.append("[SANITIZE_BUDGET_EXCEEDED]") + truncated_any = True + break + val, trunc = _recursive_smart_truncate( + i, max_len, seen, depth + 1, budget + ) if trunc: truncated_any = True new_list.append(val) @@ -339,23 +402,31 @@ def _recursive_smart_truncate( elif dataclasses.is_dataclass(obj) and not isinstance(obj, type): # Manually iterate fields to preserve 'seen' context, avoiding dataclasses.asdict recursion as_dict = {f.name: getattr(obj, f.name) for f in dataclasses.fields(obj)} - return _recursive_smart_truncate(as_dict, max_len, seen) + return _recursive_smart_truncate( + as_dict, max_len, seen, depth + 1, budget + ) elif hasattr(obj, "model_dump") and callable(obj.model_dump): # Pydantic v2 try: - return _recursive_smart_truncate(obj.model_dump(), max_len, seen) + return _recursive_smart_truncate( + obj.model_dump(), max_len, seen, depth + 1, budget + ) except Exception: pass elif hasattr(obj, "dict") and callable(obj.dict): # Pydantic v1 try: - return _recursive_smart_truncate(obj.dict(), max_len, seen) + return _recursive_smart_truncate( + obj.dict(), max_len, seen, depth + 1, budget + ) except Exception: pass elif hasattr(obj, "to_dict") and callable(obj.to_dict): # Common pattern for custom objects try: - return _recursive_smart_truncate(obj.to_dict(), max_len, seen) + return _recursive_smart_truncate( + obj.to_dict(), max_len, seen, depth + 1, budget + ) except Exception: pass elif obj is None or isinstance(obj, (int, float, bool)): @@ -1374,10 +1445,75 @@ def _truncate(self, text: str) -> tuple[str, bool]: ) return text, False + def _sanitize_external_uri( + self, uri: Optional[str] + ) -> tuple[Optional[str], bool]: + """Removes credential-bearing components from a caller-supplied URI. + + A signed GCS or HTTP URL carries its signature in the query string, and any + URI may carry ``user:password@`` userinfo. Both are bearer credentials, so + neither belongs in a stored analytics row. The scheme, host and path are + kept, which is what identifies the object. + + Args: + uri: The URI to sanitize, or None. + + Returns: + A tuple of (sanitized_uri, content_removed). + """ + if uri is None: + return None, False + if not isinstance(uri, str): + return _REDACTED_URI, True + if len(uri) > _MAX_URI_LENGTH: + return _REDACTED_URI, True + try: + parsed = urlsplit(uri) + has_userinfo = parsed.username is not None or parsed.password is not None + except ValueError: + return _REDACTED_URI, True + if has_userinfo: + # Userinfo is a credential-bearing surface by definition; do not try to + # keep the username while guessing whether it is sensitive. + return _REDACTED_URI, True + if not parsed.query and not parsed.fragment: + return uri, False + return ( + urlunsplit((parsed.scheme, parsed.netloc, parsed.path, "", "")), + True, + ) + async def _parse_content_object( - self, content: types.Content | types.Part + self, + content: types.Content | types.Part, + *, + trace_id: Optional[str] = None, + span_id: Optional[str] = None, + parse_uid: str = "", + content_ordinal: int = 0, ) -> tuple[str, list[dict[str, Any]], bool]: - """Parses a Content or Part object into summary text and content parts.""" + """Parses a Content or Part object into summary text and content parts. + + Args: + content: The Content or Part to parse. + trace_id: Trace id of the calling event. GCS object paths are built + from this argument rather than the instance field, because the + parser is shared across concurrent events and an await inside this + method can resume under another event's identity. Falls back to + the instance field. + span_id: Span id of the calling event, with the same rationale. + parse_uid: Unique per parse() call. Disambiguates object names across + concurrent events. Generated here when not supplied. + content_ordinal: Index of this Content within the calling request. The + part index restarts at zero for each Content, so without this two + messages in one request collide on the same object name. + + Returns: + A tuple of (summary_text, content_parts, is_truncated). + """ + trace_id = trace_id if trace_id is not None else self.trace_id + span_id = span_id if span_id is not None else self.span_id + parse_uid = parse_uid or uuid.uuid4().hex content_parts = [] is_truncated = False summary_text = [] @@ -1397,14 +1533,22 @@ async def _parse_content_object( # CASE A: It is already a URI (e.g. from user input) if hasattr(part, "file_data") and part.file_data: part_data["storage_mode"] = "EXTERNAL_URI" - part_data["uri"] = part.file_data.file_uri + safe_uri, uri_content_removed = self._sanitize_external_uri( + part.file_data.file_uri + ) + part_data["uri"] = safe_uri + if uri_content_removed: + is_truncated = True part_data["mime_type"] = part.file_data.mime_type # CASE B: It is Binary/Inline Data (Image/Blob) elif hasattr(part, "inline_data") and part.inline_data: if self.offloader: ext = mimetypes.guess_extension(part.inline_data.mime_type) or ".bin" - path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}{ext}" + path = ( + f"{datetime.now().date()}/{trace_id}/{span_id}_{parse_uid}" + f"_c{content_ordinal}_p{idx}{ext}" + ) try: uri = await self.offloader.upload_content( part.inline_data.data, part.inline_data.mime_type, path @@ -1443,7 +1587,10 @@ async def _parse_content_object( if self.offloader and (exceeds_inline_byte_limit or exceeds_char_limit): # Text is too big, treat as file - path = f"{datetime.now().date()}/{self.trace_id}/{self.span_id}_p{idx}.txt" + path = ( + f"{datetime.now().date()}/{trace_id}/{span_id}_{parse_uid}" + f"_c{content_ordinal}_p{idx}.txt" + ) try: uri = await self.offloader.upload_content( part.text, "text/plain", path @@ -1491,8 +1638,32 @@ async def _parse_content_object( return summary_str, content_parts, is_truncated - async def parse(self, content: Any) -> tuple[Any, list[dict[str, Any]], bool]: - """Parses content into JSON payload and content parts, potentially offloading to GCS.""" + async def parse( + self, + content: Any, + *, + trace_id: Optional[str] = None, + span_id: Optional[str] = None, + ) -> tuple[Any, list[dict[str, Any]], bool]: + """Parses content into JSON payload and content parts, potentially offloading to GCS. + + Args: + content: The content to parse. + trace_id: Trace id of the calling event, used to build GCS object + paths. Pass it per call: the parser instance is shared across + concurrent events, so a path built from the mutable instance field + can pick up another event's identity across an await. Falls back + to the instance field. + span_id: Span id of the calling event, with the same rationale. + + Returns: + A tuple of (json_payload, content_parts, is_truncated). + """ + trace_id = trace_id if trace_id is not None else self.trace_id + span_id = span_id if span_id is not None else self.span_id + # Unique per parse() call, so two events offloading at the same time + # cannot produce the same object name. + parse_uid = uuid.uuid4().hex json_payload = {} content_parts = [] is_truncated = False @@ -1508,9 +1679,15 @@ def process_text(t: str) -> tuple[str, bool]: if isinstance(content.contents, list) else [content.contents] ) - for c in contents: + for content_idx, c in enumerate(contents): role = getattr(c, "role", "unknown") - summary, parts, trunc = await self._parse_content_object(c) + summary, parts, trunc = await self._parse_content_object( + c, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=content_idx, + ) if trunc: is_truncated = True content_parts.extend(parts) @@ -1528,14 +1705,25 @@ def process_text(t: str) -> tuple[str, bool]: is_truncated = True json_payload["system_prompt"] = truncated_si else: - summary, parts, trunc = await self._parse_content_object(si) + summary, parts, trunc = await self._parse_content_object( + si, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + content_ordinal=len(contents), + ) if trunc: is_truncated = True content_parts.extend(parts) json_payload["system_prompt"] = summary elif isinstance(content, (types.Content, types.Part)): - summary, parts, trunc = await self._parse_content_object(content) + summary, parts, trunc = await self._parse_content_object( + content, + trace_id=trace_id, + span_id=span_id, + parse_uid=parse_uid, + ) return {"text_summary": summary}, parts, trunc elif isinstance(content, (dict, list)): @@ -2886,8 +3074,17 @@ async def _log_event( if self.config.content_formatter: try: raw_content = self.config.content_formatter(raw_content, event_type) - except Exception as e: - logger.warning("Content formatter failed: %s", e) + except Exception: + # Fail closed. The formatter is the operator's redaction boundary, so + # a failure must not fall back to the unformatted payload. The message + # is constant on purpose: an exception's own text or traceback can + # quote the content the formatter exists to remove. + logger.warning( + "Content formatter failed for event %s; writing a sentinel" + " instead of the original content.", + event_type, + ) + raw_content = _FORMATTER_FAILED_SENTINEL trace_id, span_id, parent_span_id = self._resolve_ids( event_data, callback_context @@ -2897,17 +3094,28 @@ async def _log_event( logger.warning("Parser not initialized; skipping event %s.", event_type) return - # Update parser's trace/span IDs for GCS pathing (reuse instance) - self.parser.trace_id = trace_id or "no_trace" - self.parser.span_id = span_id or "no_span" + # Pass the ids per call rather than assigning them to the shared parser: + # two events in flight at once would otherwise overwrite each other's + # identity between the assignment and the offload that follows an await. content_json, content_parts, parser_truncated = await self.parser.parse( - raw_content + raw_content, + trace_id=trace_id or "no_trace", + span_id=span_id or "no_span", ) is_truncated = is_truncated or parser_truncated latency_json = self._extract_latency(event_data) attributes = self._enrich_attributes(event_data, callback_context) + # Final pass over the complete assembled tree. _enrich_attributes copies + # extra_attributes (which carries session state deltas) and custom_tags in + # untouched, so this is the only point at which every value is guaranteed + # to have seen the sensitive-key redaction. + attributes, attrs_truncated = _recursive_smart_truncate( + attributes, self.config.max_content_length + ) + is_truncated = is_truncated or attrs_truncated + # Serialize attributes to JSON string try: attributes_json = json.dumps(attributes) diff --git a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py index 5719adf2b46..4bef3c341e2 100644 --- a/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py +++ b/tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py @@ -74,6 +74,9 @@ def mock_agent(): # Mock the 'name' property type(mock_a).name = mock.PropertyMock(return_value="MyTestAgent") type(mock_a).instruction = mock.PropertyMock(return_value="Test Instruction") + # root_agent returns itself (no parent), so root_agent.name is a real name + # rather than a bare mock. + mock_a.root_agent = mock_a return mock_a @@ -426,6 +429,105 @@ def test_recursive_smart_truncate_redaction(): assert truncated["nested"]["normal"] == "value" +def test_recursive_smart_truncate_bounds_self_generating_objects(): + """An object that makes a fresh wrapper on each access stops at the cap.""" + + class Endless: + """to_dict() hands back a new object every time, so ids never repeat.""" + + def to_dict(self): + return {"next": Endless()} + + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate( + {"root": Endless()}, 1000 + ) + ) + + assert is_truncated + flattened = json.dumps(truncated) + assert "[MAX_DEPTH_EXCEEDED]" in flattened + + +def test_recursive_smart_truncate_bounds_branching_self_generating_objects(): + """The depth cap alone does not bound an object that branches. + + Every access hands back two fresh children, so ids never repeat and the + 50-level cap still leaves tens of millions of nodes below it. The node + budget is what stops the walk. + """ + max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES + # Stop manufacturing children well past the budget, so a walk that is not + # bounded fails this test in a second or two rather than the minute-plus + # it takes to exhaust the depth cap on its own. + safety_limit = 5 * max_nodes + visited = 0 + + class Branching: + """to_dict() hands back two new objects every time.""" + + def to_dict(self): + nonlocal visited + visited += 1 + if visited > safety_limit: + return {"left": "stopped", "right": "stopped"} + return {"left": Branching(), "right": Branching()} + + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate( + {"root": Branching()}, 1000 + ) + ) + + assert is_truncated + assert visited <= max_nodes + assert "[SANITIZE_BUDGET_EXCEEDED]" in json.dumps(truncated) + + +def test_recursive_smart_truncate_elides_the_remainder_of_a_wide_value(): + """A wide value stops at the budget and leaves one remainder sentinel.""" + max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES + wide = list(range(max_nodes * 2)) + + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate( + {"wide": wide}, 1000 + ) + ) + + assert is_truncated + assert truncated["wide"][-1] == "[SANITIZE_BUDGET_EXCEEDED]" + assert truncated["wide"].count("[SANITIZE_BUDGET_EXCEEDED]") == 1 + # Bounded output: budget entries plus the single remainder sentinel. + assert len(truncated["wide"]) <= max_nodes + 1 + + +def test_recursive_smart_truncate_charges_directly_redacted_keys(): + """Redacted keys cost budget, so a wide temp: mapping cannot bypass it.""" + max_nodes = bigquery_agent_analytics_plugin._MAX_SANITIZE_NODES + wide_temp = {f"temp:{i}": i for i in range(max_nodes * 2)} + + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate(wide_temp, 1000) + ) + + assert is_truncated + assert len(truncated) <= max_nodes + 1 + assert "[SANITIZE_BUDGET_EXCEEDED]" in truncated + + +def test_recursive_smart_truncate_keeps_ordinary_nesting(): + """Nesting well inside the cap is copied through untouched.""" + obj = {"a": {"b": {"c": {"d": "leaf"}}}} + + truncated, is_truncated = ( + bigquery_agent_analytics_plugin._recursive_smart_truncate(obj, 1000) + ) + + assert not is_truncated + assert truncated == obj + + class TestBigQueryAgentAnalyticsPlugin: """Tests for the BigQueryAgentAnalyticsPlugin.""" @@ -685,7 +787,7 @@ async def test_content_formatter_error( dummy_arrow_schema, mock_asyncio_to_thread, ): - """Test content formatter error handling.""" + """A formatter that raises must not let the original content through.""" _ = mock_auth_default _ = mock_bq_client @@ -710,8 +812,9 @@ def error_formatter(content, event_type): log_entry = await _get_captured_event_dict_async( mock_write_client, dummy_arrow_schema ) - # If formatter fails, it logs a warning and continues with original content. - assert log_entry["content"] == '{"text_summary": "Secret message"}' + # The formatter is the redaction boundary, so its failure fails closed: + # a sentinel is written and the payload never reaches the row. + assert log_entry["content"] == "[FORMATTER_FAILED]" @pytest.mark.asyncio async def test_max_content_length( @@ -1753,6 +1856,44 @@ async def test_log_event_with_custom_tags( attributes = json.loads(log_entry["attributes"]) assert attributes["custom_tags"] == custom_tags + @pytest.mark.asyncio + async def test_custom_tags_and_extra_attributes_are_redacted( + self, + bq_plugin_inst, + mock_write_client, + callback_context, + dummy_arrow_schema, + ): + """Sensitive keys are redacted wherever they sit in the attributes tree.""" + bq_plugin_inst.config.custom_tags = { + "env": "prod", + "api_key": "sk-live-should-not-be-stored", + } + + await bq_plugin_inst._log_event( + "TEST_EVENT", + callback_context, + raw_content="test content", + event_data=bigquery_agent_analytics_plugin.EventData( + extra_attributes={ + "tool": "search", + "nested": {"refresh_token": "rt-should-not-be-stored"}, + } + ), + ) + await asyncio.sleep(0.01) + log_entry = await _get_captured_event_dict_async( + mock_write_client, dummy_arrow_schema + ) + + attributes_json = log_entry["attributes"] + assert "should-not-be-stored" not in attributes_json + attributes = json.loads(attributes_json) + assert attributes["custom_tags"]["api_key"] == "[REDACTED]" + assert attributes["custom_tags"]["env"] == "prod" + assert attributes["nested"]["refresh_token"] == "[REDACTED]" + assert attributes["tool"] == "search" + @pytest.mark.asyncio async def test_on_model_error_callback_logs_correctly( self, @@ -2778,27 +2919,34 @@ async def test_parser_instance_is_reused( assert bq_plugin_inst.parser is parser_after_init @pytest.mark.asyncio - async def test_parser_trace_id_updated_per_call( + async def test_parser_ids_are_not_mutated_per_call( self, bq_plugin_inst, mock_write_client, invocation_context, dummy_arrow_schema, ): - """trace_id and span_id on the parser should update per _log_event.""" + """_log_event passes the ids per call instead of writing them on the + + shared parser. + """ parser = bq_plugin_inst.parser original_trace_id = parser.trace_id + original_span_id = parser.span_id - bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) - await bq_plugin_inst.on_user_message_callback( - invocation_context=invocation_context, - user_message=types.Content(parts=[types.Part(text="Test")]), - ) - await asyncio.sleep(0.01) + with mock.patch.object(parser, "parse", wraps=parser.parse) as mock_parse: + bigquery_agent_analytics_plugin.TraceManager.push_span(invocation_context) + await bq_plugin_inst.on_user_message_callback( + invocation_context=invocation_context, + user_message=types.Content(parts=[types.Part(text="Test")]), + ) + await asyncio.sleep(0.01) - # After logging, trace_id/span_id should have been updated - # (they're derived from TraceManager, not the initial empty strings) - assert parser.span_id != "" + assert parser.trace_id == original_trace_id + assert parser.span_id == original_span_id + _, kwargs = mock_parse.call_args + assert kwargs["span_id"] != "" + assert kwargs["span_id"] != original_span_id @pytest.mark.asyncio async def test_parser_not_recreated_with_constructor( @@ -7595,6 +7743,179 @@ async def test_no_offloader_falls_back_to_truncate(self): assert "TRUNCATED" in parts[0]["text"] +# ================================================================ +# TEST CLASS: external URI sanitization +# ================================================================ +class TestExternalUriSanitization: + """Tests that credentials are removed from stored external URIs.""" + + def _parser(self): + return bigquery_agent_analytics_plugin.HybridContentParser( + offloader=None, + trace_id="t", + span_id="s", + max_length=-1, + ) + + @pytest.mark.asyncio + async def test_signed_gcs_uri_loses_its_signature(self): + """A signed GCS URL is stored without its query string.""" + signed = ( + "https://storage.googleapis.com/bucket/report.pdf" + "?X-Goog-Algorithm=GOOG4-RSA-SHA256" + "&X-Goog-Signature=deadbeefcafe" + ) + content = types.Content( + parts=[ + types.Part( + file_data=types.FileData( + file_uri=signed, mime_type="application/pdf" + ) + ) + ] + ) + + _, parts, is_truncated = await self._parser()._parse_content_object(content) + + assert parts[0]["storage_mode"] == "EXTERNAL_URI" + assert parts[0]["uri"] == "https://storage.googleapis.com/bucket/report.pdf" + assert "X-Goog-Signature" not in parts[0]["uri"] + assert "deadbeefcafe" not in parts[0]["uri"] + assert is_truncated + + @pytest.mark.asyncio + async def test_uri_with_userinfo_is_replaced(self): + """A URI carrying user:password@ is replaced wholesale.""" + content = types.Content( + parts=[ + types.Part( + file_data=types.FileData( + file_uri="https://alice:hunter2@example.com/a.png", + mime_type="image/png", + ) + ) + ] + ) + + _, parts, _ = await self._parser()._parse_content_object(content) + + assert parts[0]["uri"] == "[REDACTED_SENSITIVE_URI]" + + @pytest.mark.asyncio + async def test_plain_uri_is_unchanged(self): + """A URI with no query, fragment or userinfo is stored verbatim.""" + content = types.Content( + parts=[ + types.Part( + file_data=types.FileData( + file_uri="gs://bucket/photo.png", mime_type="image/png" + ) + ) + ] + ) + + _, parts, is_truncated = await self._parser()._parse_content_object(content) + + assert parts[0]["uri"] == "gs://bucket/photo.png" + assert not is_truncated + + @pytest.mark.asyncio + async def test_missing_uri_stays_none(self): + """file_data with no file_uri still records None, not a sentinel.""" + content = types.Content( + parts=[types.Part(file_data=types.FileData(mime_type="image/png"))] + ) + + _, parts, _ = await self._parser()._parse_content_object(content) + + assert parts[0]["uri"] is None + + def test_overlong_uri_is_replaced(self): + """A URI beyond the inspection bound is not parsed at all.""" + long_uri = "https://example.com/" + "a" * 9000 + + safe_uri, removed = self._parser()._sanitize_external_uri(long_uri) + + assert safe_uri == "[REDACTED_SENSITIVE_URI]" + assert removed + + +# ================================================================ +# TEST CLASS: GCS offload path identity +# ================================================================ +class TestOffloadPathIdentity: + """Tests that offload paths come from the call, not shared parser state.""" + + @pytest.mark.asyncio + async def test_concurrent_parses_keep_their_own_identity(self): + """Two parses in flight at once do not write under each other's prefix.""" + paths = [] + first_upload_started = asyncio.Event() + + async def upload_content(data, mime_type, path): + paths.append(path) + if len(paths) == 1: + # Hold the first upload open until the second has begun, so both + # parses are suspended inside _parse_content_object at once. + first_upload_started.set() + await asyncio.sleep(0.05) + return f"gs://bucket/{path}" + + offloader = mock.MagicMock() + offloader.upload_content = upload_content + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=offloader, + trace_id="", + span_id="", + max_length=10, + ) + content = types.Content(parts=[types.Part(text="X" * 200)]) + + async def parse_as(trace_id, span_id): + return await parser.parse(content, trace_id=trace_id, span_id=span_id) + + task_a = asyncio.create_task(parse_as("trace-a", "span-a")) + await asyncio.wait_for(first_upload_started.wait(), timeout=5) + task_b = asyncio.create_task(parse_as("trace-b", "span-b")) + await asyncio.gather(task_a, task_b) + + assert len(paths) == 2 + assert sum("trace-a/span-a" in p for p in paths) == 1 + assert sum("trace-b/span-b" in p for p in paths) == 1 + assert parser.trace_id == "" + assert parser.span_id == "" + + @pytest.mark.asyncio + async def test_same_part_index_in_two_messages_does_not_collide(self): + """Two messages in one request get distinct object names.""" + paths = [] + + async def upload_content(data, mime_type, path): + paths.append(path) + return f"gs://bucket/{path}" + + offloader = mock.MagicMock() + offloader.upload_content = upload_content + parser = bigquery_agent_analytics_plugin.HybridContentParser( + offloader=offloader, + trace_id="t", + span_id="s", + max_length=10, + ) + llm_request = llm_request_lib.LlmRequest( + model="gemini-pro", + contents=[ + types.Content(parts=[types.Part(text="A" * 200)]), + types.Content(parts=[types.Part(text="B" * 200)]), + ], + ) + + await parser.parse(llm_request, trace_id="t1", span_id="s1") + + assert len(paths) == 2 + assert paths[0] != paths[1] + + # ================================================================ # TEST CLASS: AGENT_RESPONSE logging (Issue #87) # ================================================================