Skip to content

Commit 79fcddb

Browse files
xuanyang15copybara-github
authored andcommitted
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 <xygoogle@google.com> PiperOrigin-RevId: 856067979
1 parent 8973618 commit 79fcddb

File tree

3 files changed

+264
-4
lines changed

3 files changed

+264
-4
lines changed

src/google/adk/cli/cli_tools_click.py

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,8 @@
3636
from . import cli_deploy
3737
from .. import version
3838
from ..evaluation.constants import MISSING_EVAL_DEPENDENCIES_MESSAGE
39+
from ..features import FeatureName
40+
from ..features import override_feature_enabled
3941
from .cli import run_cli
4042
from .fast_api import get_fast_api_app
4143
from .utils import envs
@@ -48,6 +50,56 @@
4850
)
4951

5052

53+
def _apply_feature_overrides(enable_features: tuple[str, ...]) -> None:
54+
"""Apply feature overrides from CLI flags.
55+
56+
Args:
57+
enable_features: Tuple of feature names to enable.
58+
"""
59+
for features_str in enable_features:
60+
for feature_name_str in features_str.split(","):
61+
feature_name_str = feature_name_str.strip()
62+
if not feature_name_str:
63+
continue
64+
try:
65+
feature_name = FeatureName(feature_name_str)
66+
override_feature_enabled(feature_name, True)
67+
except ValueError:
68+
valid_names = ", ".join(f.value for f in FeatureName)
69+
click.secho(
70+
f"WARNING: Unknown feature name '{feature_name_str}'. "
71+
f"Valid names are: {valid_names}",
72+
fg="yellow",
73+
err=True,
74+
)
75+
76+
77+
def feature_options():
78+
"""Decorator to add feature override options to click commands."""
79+
80+
def decorator(func):
81+
@click.option(
82+
"--enable_features",
83+
help=(
84+
"Optional. Comma-separated list of feature names to enable. "
85+
"This provides an alternative to environment variables for "
86+
"enabling experimental features. Example: "
87+
"--enable_features=JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING"
88+
),
89+
multiple=True,
90+
)
91+
@functools.wraps(func)
92+
def wrapper(*args, **kwargs):
93+
enable_features = kwargs.pop("enable_features", ())
94+
if enable_features:
95+
_apply_feature_overrides(enable_features)
96+
return func(*args, **kwargs)
97+
98+
return wrapper
99+
100+
return decorator
101+
102+
51103
class HelpfulCommand(click.Command):
52104
"""Command that shows full help on error instead of just the error message.
53105
@@ -451,6 +503,7 @@ def wrapper(*args, **kwargs):
451503

452504

453505
@main.command("run", cls=HelpfulCommand)
506+
@feature_options()
454507
@adk_services_options(default_use_local_storage=True)
455508
@click.option(
456509
"--save_session",
@@ -576,6 +629,7 @@ def wrapper(*args, **kwargs):
576629

577630

578631
@main.command("eval", cls=HelpfulCommand)
632+
@feature_options()
579633
@click.argument(
580634
"agent_module_file_path",
581635
type=click.Path(
@@ -1141,6 +1195,7 @@ def wrapper(ctx, *args, **kwargs):
11411195

11421196

11431197
@main.command("web")
1198+
@feature_options()
11441199
@fast_api_common_options()
11451200
@web_options()
11461201
@adk_services_options(default_use_local_storage=True)
@@ -1243,6 +1298,7 @@ async def _lifespan(app: FastAPI):
12431298

12441299

12451300
@main.command("api_server")
1301+
@feature_options()
12461302
# The directory of agents, where each sub-directory is a single agent.
12471303
# By default, it is the current working directory
12481304
@click.argument(
Lines changed: 197 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,197 @@
1+
# Copyright 2025 Google LLC
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License");
4+
# you may not use this file except in compliance with the License.
5+
# You may obtain a copy of the License at
6+
#
7+
# http://www.apache.org/licenses/LICENSE-2.0
8+
#
9+
# Unless required by applicable law or agreed to in writing, software
10+
# distributed under the License is distributed on an "AS IS" BASIS,
11+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
# See the License for the specific language governing permissions and
13+
# limitations under the License.
14+
15+
"""Unit tests for --enable_features CLI option."""
16+
17+
from __future__ import annotations
18+
19+
import click
20+
from click.testing import CliRunner
21+
from google.adk.cli.cli_tools_click import _apply_feature_overrides
22+
from google.adk.cli.cli_tools_click import feature_options
23+
from google.adk.features._feature_registry import _FEATURE_OVERRIDES
24+
from google.adk.features._feature_registry import _WARNED_FEATURES
25+
from google.adk.features._feature_registry import FeatureName
26+
from google.adk.features._feature_registry import is_feature_enabled
27+
import pytest
28+
29+
30+
@pytest.fixture(autouse=True)
31+
def reset_feature_overrides():
32+
"""Reset feature overrides and warnings before/after each test."""
33+
_FEATURE_OVERRIDES.clear()
34+
_WARNED_FEATURES.clear()
35+
yield
36+
_FEATURE_OVERRIDES.clear()
37+
_WARNED_FEATURES.clear()
38+
39+
40+
class TestApplyFeatureOverrides:
41+
"""Tests for _apply_feature_overrides helper function."""
42+
43+
def test_single_feature(self):
44+
"""Single feature name is applied correctly."""
45+
_apply_feature_overrides(("JSON_SCHEMA_FOR_FUNC_DECL",))
46+
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
47+
48+
def test_comma_separated_features(self):
49+
"""Comma-separated feature names are applied correctly."""
50+
_apply_feature_overrides((
51+
"JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING",
52+
))
53+
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
54+
assert is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
55+
56+
def test_multiple_flag_values(self):
57+
"""Multiple --enable_features flags are applied correctly."""
58+
_apply_feature_overrides((
59+
"JSON_SCHEMA_FOR_FUNC_DECL",
60+
"PROGRESSIVE_SSE_STREAMING",
61+
))
62+
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
63+
assert is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
64+
65+
def test_whitespace_handling(self):
66+
"""Whitespace around feature names is stripped."""
67+
_apply_feature_overrides((" JSON_SCHEMA_FOR_FUNC_DECL , COMPUTER_USE ",))
68+
assert is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
69+
assert is_feature_enabled(FeatureName.COMPUTER_USE)
70+
71+
def test_empty_string_ignored(self):
72+
"""Empty strings in the list are ignored."""
73+
_apply_feature_overrides(("",))
74+
# No error should be raised
75+
76+
def test_unknown_feature_warns(self, capsys):
77+
"""Unknown feature names emit a warning."""
78+
_apply_feature_overrides(("UNKNOWN_FEATURE_XYZ",))
79+
captured = capsys.readouterr()
80+
assert "WARNING" in captured.err
81+
assert "UNKNOWN_FEATURE_XYZ" in captured.err
82+
assert "Valid names are:" in captured.err
83+
84+
85+
class TestFeatureOptionsDecorator:
86+
"""Tests for feature_options decorator."""
87+
88+
def test_decorator_adds_enable_features_option(self):
89+
"""Decorator adds --enable_features option to command."""
90+
91+
@click.command()
92+
@feature_options()
93+
def test_cmd():
94+
pass
95+
96+
runner = CliRunner()
97+
result = runner.invoke(test_cmd, ["--help"])
98+
assert "--enable_features" in result.output
99+
100+
def test_enable_features_applied_before_command(self):
101+
"""Features are enabled before the command function runs."""
102+
feature_was_enabled = []
103+
104+
@click.command()
105+
@feature_options()
106+
def test_cmd():
107+
feature_was_enabled.append(
108+
is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
109+
)
110+
111+
runner = CliRunner()
112+
runner.invoke(
113+
test_cmd,
114+
["--enable_features=JSON_SCHEMA_FOR_FUNC_DECL"],
115+
catch_exceptions=False,
116+
)
117+
assert feature_was_enabled == [True]
118+
119+
def test_multiple_enable_features_flags(self):
120+
"""Multiple --enable_features flags work correctly."""
121+
enabled_features = []
122+
123+
@click.command()
124+
@feature_options()
125+
def test_cmd():
126+
enabled_features.append(
127+
is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
128+
)
129+
enabled_features.append(
130+
is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
131+
)
132+
133+
runner = CliRunner()
134+
runner.invoke(
135+
test_cmd,
136+
[
137+
"--enable_features=JSON_SCHEMA_FOR_FUNC_DECL",
138+
"--enable_features=PROGRESSIVE_SSE_STREAMING",
139+
],
140+
catch_exceptions=False,
141+
)
142+
assert enabled_features == [True, True]
143+
144+
def test_comma_separated_enable_features(self):
145+
"""Comma-separated feature names work correctly."""
146+
enabled_features = []
147+
148+
@click.command()
149+
@feature_options()
150+
def test_cmd():
151+
enabled_features.append(
152+
is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
153+
)
154+
enabled_features.append(
155+
is_feature_enabled(FeatureName.PROGRESSIVE_SSE_STREAMING)
156+
)
157+
158+
runner = CliRunner()
159+
runner.invoke(
160+
test_cmd,
161+
[
162+
"--enable_features=JSON_SCHEMA_FOR_FUNC_DECL,PROGRESSIVE_SSE_STREAMING"
163+
],
164+
catch_exceptions=False,
165+
)
166+
assert enabled_features == [True, True]
167+
168+
def test_no_enable_features_flag(self):
169+
"""Command works without --enable_features flag."""
170+
enabled_features = []
171+
172+
@click.command()
173+
@feature_options()
174+
def test_cmd():
175+
enabled_features.append(
176+
is_feature_enabled(FeatureName.JSON_SCHEMA_FOR_FUNC_DECL)
177+
)
178+
179+
runner = CliRunner()
180+
result = runner.invoke(test_cmd, [], catch_exceptions=False)
181+
assert result.exit_code == 0
182+
assert enabled_features == [False]
183+
184+
def test_preserves_function_metadata(self):
185+
"""Decorator preserves the wrapped function's metadata."""
186+
187+
@click.command()
188+
@feature_options()
189+
def my_test_command():
190+
"""My docstring."""
191+
pass
192+
193+
# The callback should have preserved metadata
194+
assert (
195+
"my_test_command" in my_test_command.name
196+
or my_test_command.callback.__name__ == "my_test_command"
197+
)

tests/unittests/cli/test_cli_tools_click_option_mismatch.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -94,15 +94,19 @@ def test_adk_run():
9494
run_command = _get_command_by_name(main.commands, "run")
9595

9696
assert run_command is not None, "Run command not found"
97-
_check_options_in_parameters(run_command, cli_run.callback, "run")
97+
_check_options_in_parameters(
98+
run_command, cli_run.callback, "run", ignore_params={"enable_features"}
99+
)
98100

99101

100102
def test_adk_eval():
101103
"""Test that cli_eval has all required parameters."""
102104
eval_command = _get_command_by_name(main.commands, "eval")
103105

104106
assert eval_command is not None, "Eval command not found"
105-
_check_options_in_parameters(eval_command, cli_eval.callback, "eval")
107+
_check_options_in_parameters(
108+
eval_command, cli_eval.callback, "eval", ignore_params={"enable_features"}
109+
)
106110

107111

108112
def test_adk_web():
@@ -111,7 +115,10 @@ def test_adk_web():
111115

112116
assert web_command is not None, "Web command not found"
113117
_check_options_in_parameters(
114-
web_command, cli_web.callback, "web", ignore_params={"verbose"}
118+
web_command,
119+
cli_web.callback,
120+
"web",
121+
ignore_params={"verbose", "enable_features"},
115122
)
116123

117124

@@ -124,7 +131,7 @@ def test_adk_api_server():
124131
api_server_command,
125132
cli_api_server.callback,
126133
"api_server",
127-
ignore_params={"verbose"},
134+
ignore_params={"verbose", "enable_features"},
128135
)
129136

130137

0 commit comments

Comments
 (0)