From 53b67ce6340f3f3f8c3d732f9f7811e445c76359 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Tue, 20 Jan 2026 10:45:08 -0800 Subject: [PATCH] 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 PiperOrigin-RevId: 858659818 --- src/google/adk/cli/cli_tools_click.py | 62 +++++--- .../unittests/cli/test_cli_feature_options.py | 136 ++++++++++++++++-- .../test_cli_tools_click_option_mismatch.py | 14 +- 3 files changed, 179 insertions(+), 33 deletions(-) diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index c80115be..a4b4418f 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -50,28 +50,44 @@ LOG_LEVELS = click.Choice( ) -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 feature_options(): ), 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 diff --git a/tests/unittests/cli/test_cli_feature_options.py b/tests/unittests/cli/test_cli_feature_options.py index 70bfec2d..8a507f5c 100644 --- a/tests/unittests/cli/test_cli_feature_options.py +++ b/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 @@ class TestFeatureOptionsDecorator: "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] diff --git a/tests/unittests/cli/test_cli_tools_click_option_mismatch.py b/tests/unittests/cli/test_cli_tools_click_option_mismatch.py index 3c67e9ae..9b813453 100644 --- a/tests/unittests/cli/test_cli_tools_click_option_mismatch.py +++ b/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"}, )