feat: Add --disable_features CLI option to ADK CLI

This flag can be used to override default feature enable state.

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

Author: xuanyang15

src/google/adk/cli/cli_tools_click.py

@@ -50,28 +50,44 @@
 )
 
 
-def _apply_feature_overrides(enable_features: tuple[str, ...]) -> None:
+def _apply_feature_overrides(
+    *,
+    enable_features: tuple[str, ...] = (),
+    disable_features: tuple[str, ...] = (),
+) -> None:
   """Apply feature overrides from CLI flags.
 
   Args:
     enable_features: Tuple of feature names to enable.
+    disable_features: Tuple of feature names to disable.
   """
+  feature_overrides: dict[str, bool] = {}
+
   for features_str in enable_features:
     for feature_name_str in features_str.split(","):
       feature_name_str = feature_name_str.strip()
-      if not feature_name_str:
-        continue
-      try:
-        feature_name = FeatureName(feature_name_str)
-        override_feature_enabled(feature_name, True)
-      except ValueError:
-        valid_names = ", ".join(f.value for f in FeatureName)
-        click.secho(
-            f"WARNING: Unknown feature name '{feature_name_str}'. "
-            f"Valid names are: {valid_names}",
-            fg="yellow",
-            err=True,
-        )
+      if feature_name_str:
+        feature_overrides[feature_name_str] = True
+
+  for features_str in disable_features:
+    for feature_name_str in features_str.split(","):
+      feature_name_str = feature_name_str.strip()
+      if feature_name_str:
+        feature_overrides[feature_name_str] = False
+
+  # Apply all overrides
+  for feature_name_str, enabled in feature_overrides.items():
+    try:
+      feature_name = FeatureName(feature_name_str)
+      override_feature_enabled(feature_name, enabled)
+    except ValueError:
+      valid_names = ", ".join(f.value for f in FeatureName)
+      click.secho(
+          f"WARNING: Unknown feature name '{feature_name_str}'. "
+          f"Valid names are: {valid_names}",
+          fg="yellow",
+          err=True,
+      )
 
 
 def feature_options():
@@ -88,11 +104,25 @@ def decorator(func):
         ),
         multiple=True,
     )
+    @click.option(
+        "--disable_features",
+        help=(
+            "Optional. Comma-separated list of feature names to disable. "
+            "This provides an alternative to environment variables for "
+            "disabling features. Example: "
+            "--disable_features=JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING"
+        ),
+        multiple=True,
+    )
     @functools.wraps(func)
     def wrapper(*args, **kwargs):
       enable_features = kwargs.pop("enable_features", ())
-      if enable_features:
-        _apply_feature_overrides(enable_features)
+      disable_features = kwargs.pop("disable_features", ())
+      if enable_features or disable_features:
+        _apply_feature_overrides(
+            enable_features=enable_features,
+            disable_features=disable_features,
+        )
       return func(*args, **kwargs)
 
     return wrapper

tests/unittests/cli/test_cli_feature_options.py

@@ -12,8 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""Unit tests for --enable_features CLI option."""
-
 from __future__ import annotations
 
 import click
@@ -42,45 +40,96 @@ class TestApplyFeatureOverrides:
 
   def test_single_feature(self):
     """Single feature name is applied correctly."""
-    _apply_feature_overrides(("JSON_SCHEMA_FOR_FUNC_DECL",))
+    _apply_feature_overrides(enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
     assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
 
   def test_comma_separated_features(self):
     """Comma-separated feature names are applied correctly."""
-    _apply_feature_overrides((
-        "JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",
-    ))
+    _apply_feature_overrides(
+        enable_features=("JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",)
+    )
     assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
     assert is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
 
   def test_multiple_flag_values(self):
     """Multiple --enable_features flags are applied correctly."""
-    _apply_feature_overrides((
-        "JSON_SCHEMA_FOR_FUNC_DECL",
-        "PROGRESSIVE_SSE_STREAMING",
-    ))
+    _apply_feature_overrides(
+        enable_features=(
+            "JSON_SCHEMA_FOR_FUNC_DECL",
+            "PROGRESSIVE_SSE_STREAMING",
+        )
+    )
     assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
     assert is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
 
   def test_whitespace_handling(self):
     """Whitespace around feature names is stripped."""
-    _apply_feature_overrides((" JSON_SCHEMA_FOR_FUNC_DECL , COMPUTER_USE ",))
+    _apply_feature_overrides(
+        enable_features=(" JSON_SCHEMA_FOR_FUNC_DECL , COMPUTER_USE ",)
+    )
     assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
     assert is_feature_enabled(FeatureName.COMPUTER_USE)
 
   def test_empty_string_ignored(self):
     """Empty strings in the list are ignored."""
-    _apply_feature_overrides(("",))
+    _apply_feature_overrides(enable_features=("",))
     # No error should be raised
 
   def test_unknown_feature_warns(self, capsys):
     """Unknown feature names emit a warning."""
-    _apply_feature_overrides(("UNKNOWN_FEATURE_XYZ",))
+    _apply_feature_overrides(enable_features=("UNKNOWN_FEATURE_XYZ",))
     captured = capsys.readouterr()
     assert "WARNING" in captured.err
     assert "UNKNOWN_FEATURE_XYZ" in captured.err
     assert "Valid names are:" in captured.err
 
+  def test_single_disable_feature(self):
+    """Single feature name is disabled correctly."""
+    # First enable a feature
+    _apply_feature_overrides(enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
+    assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+
+    # Then disable it
+    _apply_feature_overrides(disable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
+    assert not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+
+  def test_comma_separated_disable_features(self):
+    """Comma-separated feature names are disabled correctly."""
+    # First enable features
+    _apply_feature_overrides(
+        enable_features=("JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",)
+    )
+
+    # Then disable them
+    _apply_feature_overrides(
+        disable_features=(
+            "JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",
+        )
+    )
+    assert not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+    assert not is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
+
+  def test_disable_overrides_enable(self):
+    """Disable is applied after enable, so disable wins for same feature."""
+    _apply_feature_overrides(
+        enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",),
+        disable_features=("JSON_SCHEMA_FOR_FUNC_DECL",),
+    )
+    # disable_features is processed after enable_features
+    assert not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+
+  def test_enable_and_disable_different_features(self):
+    """Enable and disable can be used together for different features."""
+    # First enable a feature that we'll disable
+    _apply_feature_overrides(enable_features=("PROGRESSIVE_SSE_STREAMING",))
+
+    _apply_feature_overrides(
+        enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",),
+        disable_features=("PROGRESSIVE_SSE_STREAMING",),
+    )
+    assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+    assert not is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
+
 
 class TestFeatureOptionsDecorator:
   """Tests for feature_options decorator."""
@@ -195,3 +244,64 @@ def my_test_command():
         "my_test_command" in my_test_command.name
         or my_test_command.callback.__name__ == "my_test_command"
     )
+
+  def test_decorator_adds_disable_features_option(self):
+    """Decorator adds --disable_features option to command."""
+
+    @click.command()
+    @feature_options()
+    def test_cmd():
+      pass
+
+    runner = CliRunner()
+    result = runner.invoke(test_cmd, ["--help"])
+    assert "--disable_features" in result.output
+
+  def test_disable_features_applied_before_command(self):
+    """Features are disabled before the command function runs."""
+    # First enable the feature via override
+    _apply_feature_overrides(enable_features=("JSON_SCHEMA_FOR_FUNC_DECL",))
+
+    feature_was_disabled = []
+
+    @click.command()
+    @feature_options()
+    def test_cmd():
+      feature_was_disabled.append(
+          not is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+      )
+
+    runner = CliRunner()
+    runner.invoke(
+        test_cmd,
+        ["--disable_features=JSON_SCHEMA_FOR_FUNC_DECL"],
+        catch_exceptions=False,
+    )
+    assert feature_was_disabled == [True]
+
+  def test_enable_and_disable_together(self):
+    """Both --enable_features and --disable_features work together."""
+    feature_states = []
+
+    @click.command()
+    @feature_options()
+    def test_cmd():
+      feature_states.append(
+          is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
+      )
+      feature_states.append(
+          is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
+      )
+
+    runner = CliRunner()
+    runner.invoke(
+        test_cmd,
+        [
+            "--enable_features=JSON_SCHEMA_FOR_FUNC_DECL",
+            "--disable_features=PROGRESSIVE_SSE_STREAMING",
+        ],
+        catch_exceptions=False,
+    )
+    # JSON_SCHEMA_FOR_FUNC_DECL should be enabled
+    # PROGRESSIVE_SSE_STREAMING should be disabled
+    assert feature_states == [True, False]

tests/unittests/cli/test_cli_tools_click_option_mismatch.py

@@ -95,7 +95,10 @@ def test_adk_run():
 
   assert run_command is not None, "Run command not found"
   _check_options_in_parameters(
-      run_command, cli_run.callback, "run", ignore_params={"enable_features"}
+      run_command,
+      cli_run.callback,
+      "run",
+      ignore_params={"enable_features", "disable_features"},
   )
 
 
@@ -105,7 +108,10 @@ def test_adk_eval():
 
   assert eval_command is not None, "Eval command not found"
   _check_options_in_parameters(
-      eval_command, cli_eval.callback, "eval", ignore_params={"enable_features"}
+      eval_command,
+      cli_eval.callback,
+      "eval",
+      ignore_params={"enable_features", "disable_features"},
   )
 
 
@@ -118,7 +124,7 @@ def test_adk_web():
       web_command,
       cli_web.callback,
       "web",
-      ignore_params={"verbose", "enable_features"},
+      ignore_params={"verbose", "enable_features", "disable_features"},
   )
 
 
@@ -131,7 +137,7 @@ def test_adk_api_server():
       api_server_command,
       cli_api_server.callback,
       "api_server",
-      ignore_params={"verbose", "enable_features"},
+      ignore_params={"verbose", "enable_features", "disable_features"},
   )