diff --git a/src/google/adk/cli/cli_tools_click.py b/src/google/adk/cli/cli_tools_click.py index b2c8299c..41e62e3d 100644 --- a/src/google/adk/cli/cli_tools_click.py +++ b/src/google/adk/cli/cli_tools_click.py @@ -293,8 +293,12 @@ class TelemetryGroup(click.Group): ) raise except BaseException as e: - exit_code = 1 - exception_type = type(e).__name__ + if isinstance(e, KeyboardInterrupt) and ctx.meta.get("server_started"): + exit_code = 0 + exception_type = "" + else: + exit_code = 1 + exception_type = type(e).__name__ raise finally: # Exclude help requests and telemetry command group itself @@ -2004,24 +2008,8 @@ def cli_web( """, fg="green", ) - try: - if ( - ctx - and read_telemetry_consent() is True - and not ctx.meta.get("telemetry_recorded") - ): - start_time = ctx.meta.get("telemetry_start_time", time.monotonic()) - collector = MetricsCollector() - collector.record_command_run( - command="web", - exit_code=0, - duration_ms=int((time.monotonic() - start_time) * 1000), - exception_type="", - ) - ctx.meta["telemetry_recorded"] = True - except Exception: # pylint: disable=broad-except - # Failsafe: telemetry errors must never crash the CLI - pass + if ctx: + ctx.meta["server_started"] = True yield # Startup is done, now app is running click.secho( """ @@ -2172,24 +2160,8 @@ def cli_api_server( @asynccontextmanager async def _lifespan(app: FastAPI) -> AsyncIterator[None]: - try: - if ( - ctx - and read_telemetry_consent() is True - and not ctx.meta.get("telemetry_recorded") - ): - start_time = ctx.meta.get("telemetry_start_time", time.monotonic()) - collector = MetricsCollector() - collector.record_command_run( - command="api_server", - exit_code=0, - duration_ms=int((time.monotonic() - start_time) * 1000), - exception_type="", - ) - ctx.meta["telemetry_recorded"] = True - except Exception: # pylint: disable=broad-except - # Failsafe: telemetry errors must never crash the CLI - pass + if ctx: + ctx.meta["server_started"] = True yield config = uvicorn.Config( diff --git a/tests/unittests/cli/utils/test_cli_tools_click.py b/tests/unittests/cli/utils/test_cli_tools_click.py index e564ba75..c16cf2f1 100644 --- a/tests/unittests/cli/utils/test_cli_tools_click.py +++ b/tests/unittests/cli/utils/test_cli_tools_click.py @@ -307,6 +307,90 @@ def test_cli_telemetry_records_early_crash( assert source["command_run"]["exception_type"] == "KeyboardInterrupt" +def test_cli_telemetry_records_clean_shutdown_on_keyboard_interrupt_after_startup( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """TelemetryGroup invoke should record clean exit on KeyboardInterrupt after server startup.""" + + monkeypatch.setattr( + "google.adk.cli.cli_tools_click.read_telemetry_consent", + lambda: True, + ) + + temp_queue = tmp_path / "telemetry_queue.jsonl" + monkeypatch.setattr( + "google.adk.cli._telemetry._constants.QUEUE_FILE", + str(temp_queue), + ) + + @click.command("dummy_web_running") + @click.pass_context + def dummy_web_running_cmd(ctx): + ctx.meta["server_started"] = True + raise KeyboardInterrupt() + + @click.group(cls=cli_tools_click.TelemetryGroup) + def test_group(): + pass + + test_group.add_command(dummy_web_running_cmd) + + runner = CliRunner() + runner.invoke(test_group, ["dummy_web_running"]) + + assert temp_queue.exists() + with open(temp_queue, "r", encoding="utf-8") as f: + lines = f.readlines() + assert len(lines) == 1 + event = json.loads(lines[0]) + source = json.loads(event["source_extension_json"]) + assert source["command_run"]["command"] == "dummy_web_running" + assert source["command_run"]["exit_code"] == 0 + assert "exception_type" not in source["command_run"] + + +def test_cli_telemetry_records_error_after_startup_on_non_interrupt( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + """TelemetryGroup invoke should record an error for non-KeyboardInterrupt exceptions after server startup.""" + + monkeypatch.setattr( + "google.adk.cli.cli_tools_click.read_telemetry_consent", + lambda: True, + ) + + temp_queue = tmp_path / "telemetry_queue.jsonl" + monkeypatch.setattr( + "google.adk.cli._telemetry._constants.QUEUE_FILE", + str(temp_queue), + ) + + @click.command("dummy_web_runtime_error") + @click.pass_context + def dummy_web_runtime_error_cmd(ctx): + ctx.meta["server_started"] = True + raise RuntimeError("Server crashed") + + @click.group(cls=cli_tools_click.TelemetryGroup) + def test_group(): + pass + + test_group.add_command(dummy_web_runtime_error_cmd) + + runner = CliRunner() + runner.invoke(test_group, ["dummy_web_runtime_error"]) + + assert temp_queue.exists() + with open(temp_queue, "r", encoding="utf-8") as f: + lines = f.readlines() + assert len(lines) == 1 + event = json.loads(lines[0]) + source = json.loads(event["source_extension_json"]) + assert source["command_run"]["command"] == "dummy_web_runtime_error" + assert source["command_run"]["exit_code"] == 1 + assert source["command_run"]["exception_type"] == "RuntimeError" + + # cli run @pytest.mark.parametrize( "cli_args,expected_session_uri,expected_artifact_uri,expected_memory_uri",