feat: Add feature decorator for the feature registry system

Co-authored-by: Xuan Yang <xygoogle@google.com>
PiperOrigin-RevId: 832503990

src/google/adk/features/__init__.py

@@ -11,3 +11,17 @@
 # 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 ._feature_decorator import experimental
+from ._feature_decorator import stable
+from ._feature_decorator import working_in_progress
+from ._feature_registry import FeatureName
+from ._feature_registry import is_feature_enabled
+
+__all__ = [
+    "experimental",
+    "stable",
+    "working_in_progress",
+    "FeatureName",
+    "is_feature_enabled",
+]

src/google/adk/features/_feature_decorator.py

@@ -0,0 +1,118 @@
+# 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 functools
+from typing import Callable
+from typing import cast
+from typing import TypeVar
+from typing import Union
+
+from ._feature_registry import _get_feature_config
+from ._feature_registry import _register_feature
+from ._feature_registry import FeatureConfig
+from ._feature_registry import FeatureName
+from ._feature_registry import FeatureStage
+from ._feature_registry import is_feature_enabled
+
+T = TypeVar("T", bound=Union[Callable, type])
+
+
+def _make_feature_decorator(
+    *,
+    feature_name: FeatureName,
+    feature_stage: FeatureStage,
+    default_on: bool = False,
+) -> Callable[[T], T]:
+  """Decorator for experimental features.
+
+  Args:
+    feature_name: The name of the feature to decorate.
+    feature_stage: The stage of the feature.
+    default_on: Whether the feature is enabled by default.
+
+  Returns:
+    A decorator that checks if the feature is enabled and raises an error if
+    not.
+  """
+  config = _get_feature_config(feature_name)
+  if config is None:
+    config = FeatureConfig(feature_stage, default_on=default_on)
+    _register_feature(feature_name, config)
+
+  if config.stage != feature_stage:
+    raise ValueError(
+        f"Feature '{feature_name}' is being defined with stage"
+        f" '{feature_stage}', but it was previously registered with stage"
+        f" '{config.stage}'. Please ensure the feature is consistently defined."
+    )
+
+  def decorator(obj: T) -> T:
+    def check_feature_enabled():
+      if not is_feature_enabled(feature_name):
+        raise RuntimeError(f"Feature {feature_name} is not enabled.")
+
+    if isinstance(obj, type):  # decorating a class
+      original_init = obj.__init__
+
+      @functools.wraps(original_init)
+      def new_init(*args, **kwargs):
+        check_feature_enabled()
+        return original_init(*args, **kwargs)
+
+      obj.__init__ = new_init
+      return cast(T, obj)
+    elif isinstance(obj, Callable):  # decorating a function
+
+      @functools.wraps(obj)
+      def wrapper(*args, **kwargs):
+        check_feature_enabled()
+        return obj(*args, **kwargs)
+
+      return cast(T, wrapper)
+
+    else:
+      raise TypeError(
+          "@experimental can only be applied to classes or callable objects"
+      )
+
+  return decorator
+
+
+def working_in_progress(feature_name: FeatureName) -> Callable[[T], T]:
+  """Decorator for working in progress features."""
+  return _make_feature_decorator(
+      feature_name=feature_name,
+      feature_stage=FeatureStage.WIP,
+      default_on=False,
+  )
+
+
+def experimental(feature_name: FeatureName) -> Callable[[T], T]:
+  """Decorator for experimental features."""
+  return _make_feature_decorator(
+      feature_name=feature_name,
+      feature_stage=FeatureStage.EXPERIMENTAL,
+      default_on=False,
+  )
+
+
+def stable(feature_name: FeatureName) -> Callable[[T], T]:
+  """Decorator for stable features."""
+  return _make_feature_decorator(
+      feature_name=feature_name,
+      feature_stage=FeatureStage.STABLE,
+      default_on=True,
+  )

src/google/adk/features/_feature_registry.py

@@ -125,8 +125,13 @@ def is_feature_enabled(feature_name: FeatureName) -> bool:
     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}"
+  feature_name_str = (
+      feature_name.value
+      if isinstance(feature_name, FeatureName)
+      else feature_name
+  )
+  enable_var = f"ADK_ENABLE_{feature_name_str}"
+  disable_var = f"ADK_DISABLE_{feature_name_str}"
   if is_env_enabled(enable_var):
     if config.stage != FeatureStage.STABLE:
       _emit_non_stable_warning_once(feature_name, config.stage)

tests/unittests/features/test_feature_decorator.py

@@ -0,0 +1,207 @@
+# 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.
+
+import os
+import warnings
+
+from google.adk.features._feature_decorator import experimental
+from google.adk.features._feature_decorator import stable
+from google.adk.features._feature_decorator import working_in_progress
+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
+import pytest
+
+
+@working_in_progress("WIP_CLASS")
+class IncompleteFeature:
+
+  def run(self):
+    return "running"
+
+
+@working_in_progress("WIP_FUNCTION")
+def wip_function():
+  return "executing"
+
+
+@experimental("EXPERIMENTAL_CLASS")
+class ExperimentalClass:
+
+  def run(self):
+    return "running"
+
+
+@experimental("EXPERIMENTAL_FUNCTION")
+def experimental_function():
+  return "executing"
+
+
+@stable("STABLE_CLASS")
+class StableClass:
+
+  def run(self):
+    return "running"
+
+
+@stable("STABLE_FUNCTION")
+def stable_function():
+  return "executing"
+
+
+@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)
+
+  # Add an existing feature to the registry
+  _register_feature(
+      "ENABLED_EXPERIMENTAL_FEATURE",
+      FeatureConfig(FeatureStage.EXPERIMENTAL, default_on=True),
+  )
+
+  _register_feature(
+      "EXPERIMENTAL_FUNCTION",
+      FeatureConfig(FeatureStage.EXPERIMENTAL, default_on=True),
+  )
+
+
+def test_working_in_progress_stage_mismatch():
+  """Test that working_in_progress is used with a non-WIP stage."""
+  try:
+
+    @working_in_progress("ENABLED_EXPERIMENTAL_FEATURE")
+    def unused_function():  # pylint: disable=unused-variable
+      return "unused"
+
+    assert False, "Expected ValueError to be raised."
+  except ValueError as e:
+    assert (
+        "Feature 'ENABLED_EXPERIMENTAL_FEATURE' is being defined with stage"
+        " 'FeatureStage.WIP', but it was previously registered with stage"
+        " 'FeatureStage.EXPERIMENTAL'."
+        in str(e)
+    )
+
+
+def test_working_in_progress_class_raises_error():
+  """Test that WIP class raises RuntimeError by default."""
+
+  try:
+    IncompleteFeature()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature WIP_CLASS is not enabled." in str(e)
+
+
+def test_working_in_progress_class_bypass_with_env_var(monkeypatch):
+  """Test that WIP class can be bypassed with env var."""
+
+  monkeypatch.setenv("ADK_ENABLE_WIP_CLASS", "true")
+
+  with warnings.catch_warnings(record=True) as w:
+    feature = IncompleteFeature()
+    feature.run()
+    assert len(w) == 1
+    assert "[WIP] feature WIP_CLASS is enabled." in str(w[0].message)
+
+
+def test_working_in_progress_function_raises_error():
+  """Test that WIP function raises RuntimeError by default."""
+
+  try:
+    wip_function()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature WIP_FUNCTION is not enabled." in str(e)
+
+
+def test_working_in_progress_function_bypass_with_env_var(monkeypatch):
+  """Test that WIP function can be bypassed with env var."""
+
+  monkeypatch.setenv("ADK_ENABLE_WIP_FUNCTION", "true")
+
+  with warnings.catch_warnings(record=True) as w:
+    wip_function()
+    assert len(w) == 1
+    assert "[WIP] feature WIP_FUNCTION is enabled." in str(w[0].message)
+
+
+def test_disabled_experimental_class_raises_error():
+  """Test that disabled experimental class raises RuntimeError by default."""
+
+  try:
+    ExperimentalClass()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature EXPERIMENTAL_CLASS is not enabled." in str(e)
+
+
+def test_disabled_experimental_class_bypass_with_env_var(monkeypatch):
+  """Test that disabled experimental class can be bypassed with env var."""
+
+  monkeypatch.setenv("ADK_ENABLE_EXPERIMENTAL_CLASS", "true")
+
+  with warnings.catch_warnings(record=True) as w:
+    feature = ExperimentalClass()
+    feature.run()
+    assert len(w) == 1
+    assert "[EXPERIMENTAL] feature EXPERIMENTAL_CLASS is enabled." in str(
+        w[0].message
+    )
+
+
+def test_enabled_experimental_function_does_not_raise_error():
+  """Test that enabled experimental function does not raise error."""
+
+  with warnings.catch_warnings(record=True) as w:
+    experimental_function()
+    assert len(w) == 1
+    assert "[EXPERIMENTAL] feature EXPERIMENTAL_FUNCTION is enabled." in str(
+        w[0].message
+    )
+
+
+def test_enabled_experimental_function_disabled_by_env_var(monkeypatch):
+  """Test that enabled experimental function can be disabled by env var."""
+
+  monkeypatch.setenv("ADK_DISABLE_EXPERIMENTAL_FUNCTION", "true")
+
+  try:
+    experimental_function()
+    assert False, "Expected RuntimeError to be raised."
+  except RuntimeError as e:
+    assert "Feature EXPERIMENTAL_FUNCTION is not enabled." in str(e)
+
+
+def test_stable_class_does_not_raise_error_or_warn():
+  """Test that stable class does not raise error or warn."""
+
+  with warnings.catch_warnings(record=True) as w:
+    StableClass().run()
+    assert not w
+
+
+def test_stable_function_does_not_raise_error_or_warn():
+  """Test that stable function does not raise error or warn."""
+
+  with warnings.catch_warnings(record=True) as w:
+    stable_function()
+    assert not w

tests/unittests/features/test_feature_registry.py

@@ -17,13 +17,13 @@ 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
+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)
@@ -44,17 +44,11 @@ def reset_env_and_registry(monkeypatch):
     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()