Address PR review feedback (#11)

- Use one loop-state lookup\n- Preserve normalized key collisions

Author: caohy1988

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -548,23 +548,45 @@ def _normalize_json_native(
       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:
-          out_dict["[SANITIZE_BUDGET_EXCEEDED]"] = "[SANITIZE_BUDGET_EXCEEDED]"
+          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()
           if k_lower in _SENSITIVE_KEYS or k_lower.startswith("temp:"):
-            budget[0] -= 1
-            out_dict[k] = "[REDACTED]"
-            continue
+            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
         )
@@ -3627,8 +3649,9 @@ async def _get_loop_state(
     if self._generation != generation:
       raise RuntimeError("BigQuery plugin is shutting down.")
     self._cleanup_stale_loop_states()
-    if loop in self._loop_state_by_loop:
-      return self._loop_state_by_loop[loop]
+    state = self._loop_state_by_loop.get(loop)
+    if state is not None:
+      return state
 
     # grpc.aio clients are loop-bound, so we create one per event loop.
 

tests/unittests/plugins/test_bigquery_agent_analytics_plugin.py

@@ -10828,6 +10828,63 @@ def test_normalizer_redacts_and_bounds(self):
       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_native_secret_mapping_via_model_field_redacted(
       self,
@@ -12822,6 +12879,30 @@ async def blocking_shutdown(timeout=None):
     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