chore: Add override_feature_enabled to override the default feature enable states

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

Author: xuanyang15

src/google/adk/features/__init__.py

@@ -17,11 +17,13 @@
 from ._feature_decorator import working_in_progress
 from ._feature_registry import FeatureName
 from ._feature_registry import is_feature_enabled
+from ._feature_registry import override_feature_enabled
 
 __all__ = [
     "experimental",
     "stable",
     "working_in_progress",
     "FeatureName",
     "is_feature_enabled",
+    "override_feature_enabled",
 ]

src/google/adk/features/_feature_registry.py

@@ -101,6 +101,9 @@ class FeatureConfig:
 # Track which experimental features have already warned (warn only once)
 _WARNED_FEATURES: set[FeatureName] = set()
 
+# Programmatic overrides (highest priority, checked before env vars)
+_FEATURE_OVERRIDES: dict[FeatureName, bool] = {}
+
 
 def _get_feature_config(
     feature_name: FeatureName,
@@ -129,12 +132,45 @@ def _register_feature(
   _FEATURE_REGISTRY[feature_name] = config
 
 
+def override_feature_enabled(
+    feature_name: FeatureName,
+    enabled: bool,
+) -> None:
+  """Programmatically override a feature's enabled state.
+
+  This override takes highest priority, superseding environment variables
+  and registry defaults. Use this when environment variables are not
+  available or practical in your deployment environment.
+
+  Args:
+    feature_name: The feature name to override.
+    enabled: Whether the feature should be enabled.
+
+  Example:
+    ```python
+    from google.adk.features import FeatureName, override_feature_enabled
+
+    # Enable a feature programmatically
+    override_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL, True)
+    ```
+  """
+  config = _get_feature_config(feature_name)
+  if config is None:
+    raise ValueError(f"Feature {feature_name} is not registered.")
+  _FEATURE_OVERRIDES[feature_name] = enabled
+
+
 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.
 
+  Priority order (highest to lowest):
+    1. Programmatic overrides (via override_feature_enabled)
+    2. Environment variables (ADK_ENABLE_* / ADK_DISABLE_*)
+    3. Registry defaults
+
   Args:
     feature_name: The feature name (e.g., FeatureName.RESUMABILITY).
 
@@ -156,7 +192,14 @@ def _execute_agent_loop():
   if config is None:
     raise ValueError(f"Feature {feature_name} is not registered.")
 
-  # Check environment variables first (highest priority)
+  # Check programmatic overrides first (highest priority)
+  if feature_name in _FEATURE_OVERRIDES:
+    enabled = _FEATURE_OVERRIDES[feature_name]
+    if enabled and config.stage != FeatureStage.STABLE:
+      _emit_non_stable_warning_once(feature_name, config.stage)
+    return enabled
+
+  # Check environment variables second
   feature_name_str = (
       feature_name.value
       if isinstance(feature_name, FeatureName)

tests/unittests/features/test_feature_registry.py

@@ -17,13 +17,15 @@
 import os
 import warnings
 
+from google.adk.features._feature_registry import _FEATURE_OVERRIDES
 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 override_feature_enabled
 import pytest
 
 FEATURE_CONFIG_WIP = FeatureConfig(FeatureStage.WIP, default_on=False)
@@ -38,7 +40,7 @@
 
 @pytest.fixture(autouse=True)
 def reset_env_and_registry(monkeypatch):
-  """Reset environment variables and registry before each test."""
+  """Reset environment variables, registry and overrides before each test."""
   # Clean up environment variables
   for key in list(os.environ.keys()):
     if key.startswith("ADK_ENABLE_") or key.startswith("ADK_DISABLE_"):
@@ -47,11 +49,17 @@ def reset_env_and_registry(monkeypatch):
   # Reset warned features set
   _WARNED_FEATURES.clear()
 
+  # Reset feature overrides
+  _FEATURE_OVERRIDES.clear()
+
   yield
 
   # Reset warned features set
   _WARNED_FEATURES.clear()
 
+  # Reset feature overrides
+  _FEATURE_OVERRIDES.clear()
+
 
 class TestGetFeatureConfig:
   """Tests for get_feature_config() function."""
@@ -159,3 +167,76 @@ def test_warn_once_per_feature(self, monkeypatch):
       assert "[EXPERIMENTAL] feature DISABLED_FEATURE is enabled." in str(
           w[0].message
       )
+
+
+class TestOverrideFeatureEnabled:
+  """Tests for override_feature_enabled() function."""
+
+  def test_override_not_in_registry_raises_value_error(self):
+    """Overriding features not in registry raises ValueError."""
+    with pytest.raises(ValueError):
+      override_feature_enabled("UNKNOWN_FEATURE", True)
+
+  def test_override_enables_disabled_feature(self):
+    """Programmatic override can enable a disabled feature."""
+    _register_feature("OVERRIDE_TEST", FEATURE_CONFIG_EXPERIMENTAL_DISABLED)
+    assert not is_feature_enabled("OVERRIDE_TEST")
+
+    override_feature_enabled("OVERRIDE_TEST", True)
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("OVERRIDE_TEST")
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] feature OVERRIDE_TEST is enabled." in str(
+          w[0].message
+      )
+
+  def test_override_disables_enabled_feature(self):
+    """Programmatic override can disable an enabled feature."""
+    _register_feature("OVERRIDE_TEST", FEATURE_CONFIG_EXPERIMENTAL_ENABLED)
+
+    override_feature_enabled("OVERRIDE_TEST", False)
+    with warnings.catch_warnings(record=True) as w:
+      assert not is_feature_enabled("OVERRIDE_TEST")
+      assert not w
+
+  def test_override_takes_precedence_over_env_enable(self, monkeypatch):
+    """Programmatic override takes precedence over ADK_ENABLE_* env var."""
+    _register_feature("PRIORITY_TEST", FEATURE_CONFIG_EXPERIMENTAL_DISABLED)
+
+    # Set env var to enable
+    monkeypatch.setenv("ADK_ENABLE_PRIORITY_TEST", "true")
+    assert is_feature_enabled("PRIORITY_TEST")
+
+    # But override to disable
+    override_feature_enabled("PRIORITY_TEST", False)
+
+    with warnings.catch_warnings(record=True) as w:
+      assert not is_feature_enabled("PRIORITY_TEST")
+      assert not w
+
+  def test_override_takes_precedence_over_env_disable(self, monkeypatch):
+    """Programmatic override takes precedence over ADK_DISABLE_* env var."""
+    _register_feature("PRIORITY_TEST", FEATURE_CONFIG_EXPERIMENTAL_ENABLED)
+
+    # Set env var to disable
+    monkeypatch.setenv("ADK_DISABLE_PRIORITY_TEST", "true")
+    assert not is_feature_enabled("PRIORITY_TEST")
+
+    # But override to enable
+    override_feature_enabled("PRIORITY_TEST", True)
+
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("PRIORITY_TEST")
+      assert len(w) == 1
+      assert "[EXPERIMENTAL] feature PRIORITY_TEST is enabled." in str(
+          w[0].message
+      )
+
+  def test_override_stable_feature_no_warning(self):
+    """Overriding stable features does not emit warnings."""
+    _register_feature("STABLE_OVERRIDE", FEATURE_CONFIG_STABLE)
+
+    override_feature_enabled("STABLE_OVERRIDE", True)
+    with warnings.catch_warnings(record=True) as w:
+      assert is_feature_enabled("STABLE_OVERRIDE")
+      assert not w