feat: Introduce a feature registry system for ADK

xuanyang15 · copybara-github · commit 23ad40bad27a · 2025-11-13T16:14:49.000-08:00
Co-authored-by: Xuan Yang &lt;xygoogle@google.com&gt;
PiperOrigin-RevId: 832050198
diff --git a/src/google/adk/features/__init__.py b/src/google/adk/features/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
diff --git a/src/google/adk/features/feature_registry.py b/src/google/adk/features/feature_registry.py
@@ -0,0 +1,158 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import Enum
+import warnings
+
+from ..utils.env_utils import is_env_enabled
+
+
+class FeatureName(str, Enum):
+  """Feature names."""
+
+  JSON_SCHEMA_FOR_FUNC_DECL = "JSON_SCHEMA_FOR_FUNC_DECL"
+  COMPUTER_USE = "COMPUTER_USE"
+
+
+class FeatureStage(Enum):
+  """Feature lifecycle stages.
+
+  Attributes:
+    WIP: Work in progress, not functioning completely. ADK internal development
+      only.
+    EXPERIMENTAL: Feature works but API may change.
+    STABLE: Production-ready, no breaking changes without MAJOR version bump.
+  """
+
+  WIP = "wip"
+  EXPERIMENTAL = "experimental"
+  STABLE = "stable"
+
+
+@dataclass
+class FeatureConfig:
+  """Feature configuration.
+
+  Attributes:
+    stage: The feature stage.
+    default_on: Whether the feature is enabled by default.
+  """
+
+  stage: FeatureStage
+  default_on: bool = False
+
+
+# Central registry: FeatureName -> FeatureConfig
+_FEATURE_REGISTRY: dict[FeatureName, FeatureConfig] = {
+    FeatureName.JSON_SCHEMA_FOR_FUNC_DECL: FeatureConfig(
+        FeatureStage.WIP, default_on=False
+    ),
+    FeatureName.COMPUTER_USE: FeatureConfig(
+        FeatureStage.EXPERIMENTAL, default_on=True
+    ),
+}
+
+# Track which experimental features have already warned (warn only once)
+_WARNED_FEATURES: set[FeatureName] = set()
+
+
+def _get_feature_config(
+    feature_name: FeatureName,
+) -> FeatureConfig | None:
+  """Get the stage of a feature from the registry.
+
+  Args:
+    feature_name: The feature name.
+
+  Returns:
+    The feature config from the registry, or None if not found.
+  """
+  return _FEATURE_REGISTRY.get(feature_name, None)
+
+
+def _register_feature(
+    feature_name: FeatureName,
+    config: FeatureConfig,
+) -> None:
+  """Register a feature with a specific config.
+
+  Args:
+    feature_name: The feature name.
+    config: The feature config to register.
+  """
+  _FEATURE_REGISTRY[feature_name] = config
+
+
+def is_feature_enabled(feature_name: FeatureName) -> bool:
+  """Check if a feature is enabled at runtime.
+
+  This function is used for runtime behavior gating within stable features.
+  It allows you to conditionally enable new behavior based on feature flags.
+
+  Args:
+    feature_name: The feature name (e.g., FeatureName.RESUMABILITY).
+
+  Returns:
+    True if the feature is enabled, False otherwise.
+
+  Example:
+    ```python
+    def _execute_agent_loop():
+      if is_feature_enabled(FeatureName.RESUMABILITY):
+        # New behavior: save checkpoints for resuming
+        return _execute_with_checkpoints()
+      else:
+        # Old behavior: run without checkpointing
+        return _execute_standard()
+    ```
+  """
+  config = _get_feature_config(feature_name)
+  if config is None:
+    raise ValueError(f"Feature {feature_name} is not registered.")
+
+  # Check environment variables first (highest priority)
+  enable_var = f"ADK_ENABLE_{feature_name}"
+  disable_var = f"ADK_DISABLE_{feature_name}"
+  if is_env_enabled(enable_var):
+    if config.stage != FeatureStage.STABLE:
+      _emit_non_stable_warning_once(feature_name, config.stage)
+    return True
+  if is_env_enabled(disable_var):
+    return False
+
+  # Fall back to registry config
+  if config.stage != FeatureStage.STABLE and config.default_on:
+    _emit_non_stable_warning_once(feature_name, config.stage)
+  return config.default_on
+
+
+def _emit_non_stable_warning_once(
+    feature_name: FeatureName,
+    feature_stage: FeatureStage,
+) -> None:
+  """Emit a warning for a non-stable feature, but only once per feature.
+
+  Args:
+    feature_name: The feature name.
+    feature_stage: The feature stage.
+  """
+  if feature_name not in _WARNED_FEATURES:
+    _WARNED_FEATURES.add(feature_name)
+    full_message = (
+        f"[{feature_stage.name.upper()}] feature {feature_name} is enabled."
+    )
+    warnings.warn(full_message, category=UserWarning, stacklevel=4)
diff --git a/tests/unittests/features/test_feature_registry.py b/tests/unittests/features/test_feature_registry.py
@@ -0,0 +1,167 @@
+# Copyright 2025 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from __future__ import annotations
+
+import os
+import warnings
+
+from google.adk.features.feature_registry import _FEATURE_REGISTRY
+from google.adk.features.feature_registry import _get_feature_config
+from google.adk.features.feature_registry import _register_feature
+from google.adk.features.feature_registry import _WARNED_FEATURES
+from google.adk.features.feature_registry import FeatureConfig
+from google.adk.features.feature_registry import FeatureStage
+from google.adk.features.feature_registry import is_feature_enabled
+import pytest
+
+FEATURE_CONFIG_WIP = FeatureConfig(FeatureStage.WIP, default_on=False)
+FEATURE_CONFIG_EXPERIMENTAL_DISABLED = FeatureConfig(
+    FeatureStage.EXPERIMENTAL, default_on=False
+)
+FEATURE_CONFIG_EXPERIMENTAL_ENABLED = FeatureConfig(
+    FeatureStage.EXPERIMENTAL, default_on=True
+)
+FEATURE_CONFIG_STABLE = FeatureConfig(FeatureStage.STABLE, default_on=True)
+
+
+@pytest.fixture(autouse=True)
+def reset_env_and_registry(monkeypatch):
+  """Reset environment variables and registry before each test."""
+  # Clean up environment variables
+  for key in list(os.environ.keys()):
+    if key.startswith("ADK_ENABLE_") or key.startswith("ADK_DISABLE_"):
+      monkeypatch.delenv(key, raising=False)
+
+  # Clear registry (but keep it as a dict for adding test entries)
+  _FEATURE_REGISTRY.clear()
+
+  # Reset warned features set
+  _WARNED_FEATURES.clear()
+
+  yield
+
+  # Clean up after test
+  _FEATURE_REGISTRY.clear()
+
+  # Reset warned features set
+  _WARNED_FEATURES.clear()
+
+
+class TestGetFeatureConfig:
+  """Tests for get_feature_config() function."""
+
+  def test_feature_in_registry(self):
+    """Returns correct config for features in registry."""
+    _register_feature("MY_FEATURE", FEATURE_CONFIG_EXPERIMENTAL_ENABLED)
+    assert (
+        _get_feature_config("MY_FEATURE") == FEATURE_CONFIG_EXPERIMENTAL_ENABLED
+    )
+
+  def test_feature_not_in_registry(self):
+    """Returns EXPERIMENTAL_DISABLED for features not in registry."""
+    assert _get_feature_config("UNKNOWN_FEATURE") is None
+
+
+class TestIsFeatureEnabled:
+  """Tests for is_feature_enabled() runtime check function."""
+
+  def test_not_in_registry_raises_value_error(self):
+    """Features not in registry raise ValueError when checked."""
+    with pytest.raises(ValueError):
+      is_feature_enabled("NEW_FEATURE")
+
+  def test_wip_feature_disabled(self):
+    """WIP features are disabled by default."""
+    _register_feature("WIP_FEATURE", FEATURE_CONFIG_WIP)
+    with warnings.catch_warnings(record=True) as w:
+      assert not is_feature_enabled("WIP_FEATURE")
+      assert not w
+
+  def test_wip_feature_enabled(self):
+    """WIP features are disabled by default."""
+    _register_feature(
+        "WIP_FEATURE", FeatureConfig(FeatureStage.WIP, default_on=True)
+    )
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("WIP_FEATURE")
+      assert len(w) == 1
+      assert "[WIP] feature WIP_FEATURE is enabled." in str(w[0].message)
+
+  def test_experimental_disabled_feature(self):
+    """Experimental disabled features are disabled."""
+    _register_feature("EXP_DISABLED", FEATURE_CONFIG_EXPERIMENTAL_DISABLED)
+    with warnings.catch_warnings(record=True) as w:
+      assert not is_feature_enabled("EXP_DISABLED")
+      assert not w
+
+  def test_experimental_enabled_feature(self):
+    """Experimental enabled features are enabled."""
+    _register_feature("EXP_ENABLED", FEATURE_CONFIG_EXPERIMENTAL_ENABLED)
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("EXP_ENABLED")
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] feature EXP_ENABLED is enabled." in str(
+          w[0].message
+      )
+
+  def test_stable_feature_enabled(self):
+    """Stable features are enabled."""
+    _register_feature("STABLE_FEATURE", FEATURE_CONFIG_STABLE)
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("STABLE_FEATURE")
+      assert not w
+
+  def test_enable_env_var_takes_precedence(self, monkeypatch):
+    """ADK_ENABLE_<FEATURE> takes precedence over registry."""
+    # Feature disabled in registry
+    _register_feature("DISABLED_FEATURE", FEATURE_CONFIG_EXPERIMENTAL_DISABLED)
+
+    # But enabled via env var
+    monkeypatch.setenv("ADK_ENABLE_DISABLED_FEATURE", "true")
+
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("DISABLED_FEATURE")
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] feature DISABLED_FEATURE is enabled." in str(
+          w[0].message
+      )
+
+  def test_disable_env_var_takes_precedence(self, monkeypatch):
+    """ADK_DISABLE_<FEATURE> takes precedence over registry."""
+    # Feature enabled in registry
+    _register_feature("ENABLED_FEATURE", FEATURE_CONFIG_STABLE)
+
+    # But disabled via env var
+    monkeypatch.setenv("ADK_DISABLE_ENABLED_FEATURE", "true")
+
+    with warnings.catch_warnings(record=True) as w:
+      assert not is_feature_enabled("ENABLED_FEATURE")
+      assert not w
+
+  def test_warn_once_per_feature(self, monkeypatch):
+    """Warn once per feature, even if being used multiple times."""
+    # Feature disabled in registry
+    _register_feature("DISABLED_FEATURE", FEATURE_CONFIG_EXPERIMENTAL_DISABLED)
+
+    # But enabled via env var
+    monkeypatch.setenv("ADK_ENABLE_DISABLED_FEATURE", "true")
+
+    with warnings.catch_warnings(record=True) as w:
+      is_feature_enabled("DISABLED_FEATURE")
+      is_feature_enabled("DISABLED_FEATURE")
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] feature DISABLED_FEATURE is enabled." in str(
+          w[0].message
+      )
