From 79fcddb39f71a4c1342e63b4d67832b3eccb2652 Mon Sep 17 00:00:00 2001 From: Xuan Yang Date: Tue, 13 Jan 2026 23:57:26 -0800 Subject: [PATCH] feat: Add `--enable_features` CLI option to ADK CLI This flag can be used to override default feature enable state. Co-authored-by: Xuan Yang PiperOrigin-RevId: 856067979 --- src/google/adk/cli/cli_tools_click.py | 56 +++++ .../unittests/cli/test_cli_feature_options.py | 197 ++++++++++++++++++ .../test_cli_tools_click_option_mismatch.py | 15 +- 3 files changed, 264 insertions(+), 4 deletions(-) create mode 100644 tests/unittests/cli/test_cli_feature_options.py diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index 4aa39dce..5d7611f2 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -36,6 +36,8 @@ from . import cli_create from . import cli_deploy from .. import version from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE +from ..features import FeatureName +from ..features import override_feature_enabled from .cli import run_cli from .fast_api import get_fast_api_app from .utils import envs @@ -48,6 +50,56 @@ LOG_LEVELS = click.Choice( ) +def _apply_feature_overrides(enable_features: tuple[str, ...]) -> None: + """Apply feature overrides from CLI flags. + + Args: + enable_features: Tuple of feature names to enable. + """ + 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, + ) + + +def feature_options(): + """Decorator to add feature override options to click commands.""" + + def decorator(func): + @click.option( + "--enable_features", + help=( + "Optional. Comma-separated list of feature names to enable. " + "This provides an alternative to environment variables for " + "enabling experimental features. Example: " + "--enable_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) + return func(*args, **kwargs) + + return wrapper + + return decorator + + class HelpfulCommand(click.Command): """Command that shows full help on error instead of just the error message. @@ -451,6 +503,7 @@ def adk_services_options(*, default_use_local_storage: bool = True): @main.command("run", cls=HelpfulCommand) +@feature_options() @adk_services_options(default_use_local_storage=True) @click.option( "--save_session", @@ -576,6 +629,7 @@ def eval_options(): @main.command("eval", cls=HelpfulCommand) +@feature_options() @click.argument( "agent_module_file_path", type=click.Path( @@ -1141,6 +1195,7 @@ def fast_api_common_options(): @main.command("web") +@feature_options() @fast_api_common_options() @web_options() @adk_services_options(default_use_local_storage=True) @@ -1243,6 +1298,7 @@ def cli_web( @main.command("api_server") +@feature_options() # The directory of agents, where each sub-directory is a single agent. # By default, it is the current working directory @click.argument( diff --git a/tests/unittests/cli/test_cli_feature_options.py b/tests/unittests/cli/test_cli_feature_options.py new file mode 100644 index 00000000..70bfec2d --- /dev/null +++ b/tests/unittests/cli/test_cli_feature_options.py @@ -0,0 +1,197 @@ +# 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. + +"""Unit tests for --enable_features CLI option.""" + +from __future__ import annotations + +import click +from click.testing import CliRunner +from google.adk.cli.cli_tools_click import _apply_feature_overrides +from google.adk.cli.cli_tools_click import feature_options +from google.adk.features._feature_registry import _FEATURE_OVERRIDES +from google.adk.features._feature_registry import _WARNED_FEATURES +from google.adk.features._feature_registry import FeatureName +from google.adk.features._feature_registry import is_feature_enabled +import pytest + + +@pytest.fixture(autouse=True) +def reset_feature_overrides(): + """Reset feature overrides and warnings before/after each test.""" + _FEATURE_OVERRIDES.clear() + _WARNED_FEATURES.clear() + yield + _FEATURE_OVERRIDES.clear() + _WARNED_FEATURES.clear() + + +class TestApplyFeatureOverrides: + """Tests for _apply_feature_overrides helper function.""" + + def test_single_feature(self): + """Single feature name is applied correctly.""" + _apply_feature_overrides(("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", + )) + 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", + )) + 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 ",)) + 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(("",)) + # No error should be raised + + def test_unknown_feature_warns(self, capsys): + """Unknown feature names emit a warning.""" + _apply_feature_overrides(("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 + + +class TestFeatureOptionsDecorator: + """Tests for feature_options decorator.""" + + def test_decorator_adds_enable_features_option(self): + """Decorator adds --enable_features option to command.""" + + @click.command() + @feature_options() + def test_cmd(): + pass + + runner = CliRunner() + result = runner.invoke(test_cmd, ["--help"]) + assert "--enable_features" in result.output + + def test_enable_features_applied_before_command(self): + """Features are enabled before the command function runs.""" + feature_was_enabled = [] + + @click.command() + @feature_options() + def test_cmd(): + feature_was_enabled.append( + is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL) + ) + + runner = CliRunner() + runner.invoke( + test_cmd, + ["--enable_features=JSON_SCHEMA_FOR_FUNC_DECL"], + catch_exceptions=False, + ) + assert feature_was_enabled == [True] + + def test_multiple_enable_features_flags(self): + """Multiple --enable_features flags work correctly.""" + enabled_features = [] + + @click.command() + @feature_options() + def test_cmd(): + enabled_features.append( + is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL) + ) + enabled_features.append( + is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING) + ) + + runner = CliRunner() + runner.invoke( + test_cmd, + [ + "--enable_features=JSON_SCHEMA_FOR_FUNC_DECL", + "--enable_features=PROGRESSIVE_SSE_STREAMING", + ], + catch_exceptions=False, + ) + assert enabled_features == [True, True] + + def test_comma_separated_enable_features(self): + """Comma-separated feature names work correctly.""" + enabled_features = [] + + @click.command() + @feature_options() + def test_cmd(): + enabled_features.append( + is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL) + ) + enabled_features.append( + is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING) + ) + + runner = CliRunner() + runner.invoke( + test_cmd, + [ + "--enable_features=JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING" + ], + catch_exceptions=False, + ) + assert enabled_features == [True, True] + + def test_no_enable_features_flag(self): + """Command works without --enable_features flag.""" + enabled_features = [] + + @click.command() + @feature_options() + def test_cmd(): + enabled_features.append( + is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL) + ) + + runner = CliRunner() + result = runner.invoke(test_cmd, [], catch_exceptions=False) + assert result.exit_code == 0 + assert enabled_features == [False] + + def test_preserves_function_metadata(self): + """Decorator preserves the wrapped function's metadata.""" + + @click.command() + @feature_options() + def my_test_command(): + """My docstring.""" + pass + + # The callback should have preserved metadata + assert ( + "my_test_command" in my_test_command.name + or my_test_command.callback.__name__ == "my_test_command" + ) 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 346fd421..3c67e9ae 100644 --- a/tests/unittests/cli/test_cli_tools_click_option_mismatch.py +++ b/tests/unittests/cli/test_cli_tools_click_option_mismatch.py @@ -94,7 +94,9 @@ def test_adk_run(): run_command = _get_command_by_name(main.commands, "run") assert run_command is not None, "Run command not found" - _check_options_in_parameters(run_command, cli_run.callback, "run") + _check_options_in_parameters( + run_command, cli_run.callback, "run", ignore_params={"enable_features"} + ) def test_adk_eval(): @@ -102,7 +104,9 @@ def test_adk_eval(): eval_command = _get_command_by_name(main.commands, "eval") assert eval_command is not None, "Eval command not found" - _check_options_in_parameters(eval_command, cli_eval.callback, "eval") + _check_options_in_parameters( + eval_command, cli_eval.callback, "eval", ignore_params={"enable_features"} + ) def test_adk_web(): @@ -111,7 +115,10 @@ def test_adk_web(): assert web_command is not None, "Web command not found" _check_options_in_parameters( - web_command, cli_web.callback, "web", ignore_params={"verbose"} + web_command, + cli_web.callback, + "web", + ignore_params={"verbose", "enable_features"}, ) @@ -124,7 +131,7 @@ def test_adk_api_server(): api_server_command, cli_api_server.callback, "api_server", - ignore_params={"verbose"}, + ignore_params={"verbose", "enable_features"}, )