feat: add express mode telemetry logging for ADK CLI onboarding
Track user choices (e.g. CREATE_EXPRESS, MANUAL_PROJECT, ABANDON) during Express Mode onboarding in ADK CLI telemetry logs. - Added express_mode_action field to CliCommandRun proto schema. - Recorded express_mode_action in MetricsCollector and forwarded from Click context metadata during command execution. - Added unit tests for express_mode_action serialization. - Updated Clearcut route test case ADK_CLI_basic.textpb. Co-authored-by: Kathy Wu <wukathy@google.com> PiperOrigin-RevId: 962278596
This commit is contained in:
committed by
Copybara-Service
parent
6930be4305
commit
4ccc6be6d4
@@ -196,6 +196,7 @@ class MetricsCollector:
|
||||
exit_code: int = 0,
|
||||
duration_ms: int = 0,
|
||||
exception_type: str = "",
|
||||
express_mode_action: str = "",
|
||||
) -> None:
|
||||
"""Records a command execution and safely appends to local disk queue."""
|
||||
with self._lock:
|
||||
@@ -218,6 +219,10 @@ class MetricsCollector:
|
||||
if exception_type:
|
||||
# Enforce string length limit on exception type name
|
||||
command_run["exception_type"] = exception_type[:_MAX_EXCEPTION_LENGTH]
|
||||
if express_mode_action:
|
||||
command_run["express_mode_action"] = express_mode_action[
|
||||
:_MAX_STRING_LENGTH
|
||||
]
|
||||
|
||||
source_extension = {
|
||||
"client_session_id": self._session_id,
|
||||
|
||||
@@ -345,6 +345,7 @@ class TelemetryGroup(click.Group):
|
||||
exit_code=exit_code,
|
||||
duration_ms=int((time.monotonic() - start_time) * 1000),
|
||||
exception_type=exception_type,
|
||||
express_mode_action=ctx.meta.get("express_mode_action", ""),
|
||||
)
|
||||
except Exception: # pylint: disable=broad-except
|
||||
# Failsafe: telemetry errors must never crash the CLI
|
||||
|
||||
@@ -155,6 +155,16 @@ def prompt_for_google_api_key(
|
||||
return google_api_key
|
||||
|
||||
|
||||
def _record_express_action(action_name: str) -> None:
|
||||
"""Records the onboarding choice for the CLI telemetry event."""
|
||||
# `Context.meta` is one dict shared by reference with every ancestor context,
|
||||
# so writing it here is enough for the root `TelemetryGroup` context to read
|
||||
# it back when the command finishes.
|
||||
ctx = click.get_current_context(silent=True)
|
||||
if ctx is not None:
|
||||
ctx.meta["express_mode_action"] = action_name
|
||||
|
||||
|
||||
def handle_login_with_google() -> VertexAIAuth | ExpressModeAuth:
|
||||
"""Handles the "Login with Google" flow."""
|
||||
if not gcp_utils.check_adc():
|
||||
@@ -177,6 +187,7 @@ def handle_login_with_google() -> VertexAIAuth | ExpressModeAuth:
|
||||
region = express_project.get("region", "us-central1")
|
||||
if project_id:
|
||||
click.secho(f"Using existing Express project: {project_id}", fg="green")
|
||||
_record_express_action("EXISTING_EXPRESS")
|
||||
return ExpressModeAuth(
|
||||
api_key=api_key, project_id=project_id, region=region
|
||||
)
|
||||
@@ -199,6 +210,7 @@ def handle_login_with_google() -> VertexAIAuth | ExpressModeAuth:
|
||||
type=click.IntRange(0, len(projects)),
|
||||
)
|
||||
if project_index == 0:
|
||||
_record_express_action("MANUAL_PROJECT")
|
||||
selected_project_id = prompt_for_google_cloud(None)
|
||||
else:
|
||||
selected_project_id = projects[project_index - 1][0]
|
||||
@@ -220,9 +232,11 @@ def handle_login_with_google() -> VertexAIAuth | ExpressModeAuth:
|
||||
)
|
||||
|
||||
if action == "3":
|
||||
_record_express_action("ABANDON")
|
||||
raise click.Abort()
|
||||
|
||||
if action == "1":
|
||||
_record_express_action("MANUAL_PROJECT")
|
||||
google_cloud_project = prompt_for_google_cloud(None)
|
||||
google_cloud_region = prompt_for_google_cloud_region(None)
|
||||
return VertexAIAuth(
|
||||
@@ -278,10 +292,12 @@ Choose region""",
|
||||
click.secho(
|
||||
"Failed to unset project. Please do it manually.", fg="red"
|
||||
)
|
||||
_record_express_action("CREATE_EXPRESS")
|
||||
return ExpressModeAuth(
|
||||
api_key=api_key, project_id=project_id, region=region
|
||||
)
|
||||
|
||||
_record_express_action("ABANDON")
|
||||
click.secho(_NOT_ELIGIBLE_MSG, fg="red")
|
||||
raise click.Abort()
|
||||
|
||||
|
||||
@@ -106,6 +106,7 @@ class CliMetricsTest(unittest.TestCase):
|
||||
exit_code=0,
|
||||
duration_ms=450,
|
||||
exception_type="",
|
||||
express_mode_action="CREATE_EXPRESS",
|
||||
)
|
||||
|
||||
# Verify it's written in queue file
|
||||
@@ -121,6 +122,9 @@ class CliMetricsTest(unittest.TestCase):
|
||||
self.assertEqual(source["command_run"]["subcommand"], "create")
|
||||
self.assertEqual(source["command_run"]["exit_code"], 0)
|
||||
self.assertEqual(source["command_run"]["duration_ms"], 450)
|
||||
self.assertEqual(
|
||||
source["command_run"]["express_mode_action"], "CREATE_EXPRESS"
|
||||
)
|
||||
self.assertEqual(
|
||||
source["command_run"]["flags"],
|
||||
["--debug", "--project", "-v", "--user"],
|
||||
|
||||
@@ -558,3 +558,105 @@ def test_get_gcp_region_from_gcloud_fail(
|
||||
),
|
||||
)
|
||||
assert _onboarding.get_gcp_region_from_gcloud() == ""
|
||||
|
||||
|
||||
# express_mode_action telemetry
|
||||
def _onboard_and_get_root_meta() -> Dict[str, Any]:
|
||||
"""Runs onboarding under a nested context and returns the *root* context meta.
|
||||
|
||||
`TelemetryGroup` reads `express_mode_action` off the root context, so the
|
||||
assertion has to be made there rather than on the subcommand context the
|
||||
onboarding code happens to run under.
|
||||
"""
|
||||
root_ctx = click.Context(click.Command("adk"))
|
||||
sub_ctx = click.Context(click.Command("create"), parent=root_ctx)
|
||||
with root_ctx, sub_ctx:
|
||||
try:
|
||||
_onboarding.handle_login_with_google()
|
||||
except click.Abort:
|
||||
pass
|
||||
return root_ctx.meta
|
||||
|
||||
|
||||
def test_express_action_recorded_for_existing_express(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Reusing an existing Express project records EXISTING_EXPRESS."""
|
||||
monkeypatch.setattr(gcp_utils, "check_adc", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
gcp_utils,
|
||||
"retrieve_express_project",
|
||||
lambda: {"api_key": "key", "project_id": "proj", "region": "us-central1"},
|
||||
)
|
||||
|
||||
meta = _onboard_and_get_root_meta()
|
||||
assert meta.get("express_mode_action") == "EXISTING_EXPRESS"
|
||||
|
||||
|
||||
def test_express_action_recorded_for_manual_project(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Entering a project ID by hand records MANUAL_PROJECT."""
|
||||
monkeypatch.setattr(gcp_utils, "check_adc", lambda: True)
|
||||
monkeypatch.setattr(gcp_utils, "retrieve_express_project", lambda: None)
|
||||
monkeypatch.setattr(gcp_utils, "list_gcp_projects", lambda limit: [])
|
||||
prompts = iter(["1", "test-proj", "us-east1"])
|
||||
monkeypatch.setattr(click, "prompt", lambda *a, **k: next(prompts))
|
||||
|
||||
meta = _onboard_and_get_root_meta()
|
||||
assert meta.get("express_mode_action") == "MANUAL_PROJECT"
|
||||
|
||||
|
||||
def test_express_action_recorded_for_manual_project_from_list(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Opting out of the project list to type an ID records MANUAL_PROJECT."""
|
||||
monkeypatch.setattr(gcp_utils, "check_adc", lambda: True)
|
||||
monkeypatch.setattr(gcp_utils, "retrieve_express_project", lambda: None)
|
||||
monkeypatch.setattr(
|
||||
gcp_utils, "list_gcp_projects", lambda limit: [("p1", "Project 1")]
|
||||
)
|
||||
prompts = iter([0, "manual-proj", "us-east1"])
|
||||
monkeypatch.setattr(click, "prompt", lambda *a, **k: next(prompts))
|
||||
|
||||
meta = _onboard_and_get_root_meta()
|
||||
assert meta.get("express_mode_action") == "MANUAL_PROJECT"
|
||||
|
||||
|
||||
def test_express_action_recorded_for_create_express(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Signing up for a new Express project records CREATE_EXPRESS."""
|
||||
monkeypatch.setattr(gcp_utils, "check_adc", lambda: True)
|
||||
monkeypatch.setattr(gcp_utils, "retrieve_express_project", lambda: None)
|
||||
monkeypatch.setattr(gcp_utils, "list_gcp_projects", lambda limit: [])
|
||||
monkeypatch.setattr(gcp_utils, "check_express_eligibility", lambda: True)
|
||||
monkeypatch.setattr(click, "confirm", lambda *a, **k: True)
|
||||
prompts = iter(["2", "1"])
|
||||
monkeypatch.setattr(click, "prompt", lambda *a, **k: next(prompts))
|
||||
monkeypatch.setattr(
|
||||
gcp_utils,
|
||||
"sign_up_express",
|
||||
lambda location="us-central1": {
|
||||
"api_key": "new-key",
|
||||
"project_id": "new-proj",
|
||||
"region": location,
|
||||
},
|
||||
)
|
||||
monkeypatch.setattr(_onboarding, "get_gcp_project_from_gcloud", lambda: "")
|
||||
|
||||
meta = _onboard_and_get_root_meta()
|
||||
assert meta.get("express_mode_action") == "CREATE_EXPRESS"
|
||||
|
||||
|
||||
def test_express_action_recorded_for_abandon(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""Choosing to abandon onboarding records ABANDON."""
|
||||
monkeypatch.setattr(gcp_utils, "check_adc", lambda: True)
|
||||
monkeypatch.setattr(gcp_utils, "retrieve_express_project", lambda: None)
|
||||
monkeypatch.setattr(gcp_utils, "list_gcp_projects", lambda limit: [])
|
||||
monkeypatch.setattr(click, "prompt", lambda *a, **k: "3")
|
||||
|
||||
meta = _onboard_and_get_root_meta()
|
||||
assert meta.get("express_mode_action") == "ABANDON"
|
||||
|
||||
@@ -36,6 +36,7 @@ from click.testing import CliRunner
|
||||
from google.adk.agents.base_agent import BaseAgent
|
||||
from google.adk.agents.run_config import StreamingMode
|
||||
from google.adk.cli import cli_tools_click
|
||||
from google.adk.cli.utils import gcp_utils
|
||||
from google.adk.evaluation.eval_case import EvalCase
|
||||
from google.adk.evaluation.eval_set import EvalSet
|
||||
from google.adk.evaluation.local_eval_set_results_manager import LocalEvalSetResultsManager
|
||||
@@ -263,6 +264,56 @@ def test_cli_telemetry_captures_subcommand_flags(
|
||||
assert "<app_name>" in source["command_run"]["flags"]
|
||||
|
||||
|
||||
def test_cli_telemetry_records_express_mode_action(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
"""An onboarding choice must reach the logged command_run.
|
||||
|
||||
This is the only test covering the hand-off as a whole: `_onboarding` writes
|
||||
the action into the Click context and `TelemetryGroup` reads it back out. A
|
||||
typo in the meta key on either side passes every other test in the suite.
|
||||
"""
|
||||
monkeypatch.setattr(
|
||||
"google.adk.cli.cli_tools_click.read_telemetry_consent",
|
||||
lambda: True,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"google.adk.cli._telemetry._metrics_collector"
|
||||
".MetricsCollector._is_rate_limited",
|
||||
lambda: True,
|
||||
)
|
||||
temp_queue = tmp_path / "telemetry_queue.jsonl"
|
||||
monkeypatch.setattr(
|
||||
"google.adk.cli._telemetry._constants.QUEUE_FILE",
|
||||
str(temp_queue),
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"google.adk.cli._telemetry._constants.TELEMETRY_SESSIONS_DIR",
|
||||
str(tmp_path / "telemetry_sessions"),
|
||||
)
|
||||
|
||||
# Drive `create` into the "3. Login with Google" branch, which finds an
|
||||
# existing Express project and records EXISTING_EXPRESS.
|
||||
monkeypatch.setattr(gcp_utils, "check_adc", lambda: True)
|
||||
monkeypatch.setattr(
|
||||
gcp_utils,
|
||||
"retrieve_express_project",
|
||||
lambda: {"api_key": "key", "project_id": "proj", "region": "us-central1"},
|
||||
)
|
||||
|
||||
runner = CliRunner()
|
||||
result = runner.invoke(
|
||||
cli_tools_click.main,
|
||||
["create", "--model", "gemini-2.0", str(tmp_path / "new_app")],
|
||||
input="3\n",
|
||||
)
|
||||
assert result.exit_code == 0
|
||||
|
||||
event = json.loads(temp_queue.read_text().splitlines()[0])
|
||||
source = json.loads(event["source_extension_json"])
|
||||
assert source["command_run"]["express_mode_action"] == "EXISTING_EXPRESS"
|
||||
|
||||
|
||||
def test_cli_telemetry_skips_when_already_recorded(
|
||||
tmp_path: Path, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
|
||||
Reference in New Issue
Block a user