From b0fd0a9fa3fe898c8a44f9623b63c0904809c55f Mon Sep 17 00:00:00 2001 From: Lorenze Jay <63378463+lorenzejay@users.noreply.github.com> Date: Wed, 9 Sep 2026 09:43:11 -0700 Subject: [PATCH] feat(telemetry): track checkpoint runtime and CLI usage (#7348) * feat(cli): track checkpoint command and TUI usage * feat(telemetry): track runtime checkpoint operations * fix(telemetry): count prune usage after argument validation * style(cli): format checkpoint prune command --------- Co-authored-by: Vidit Ostwal <110953813+Vidit-Ostwal@users.noreply.github.com> --- lib/cli/src/crewai_cli/checkpoint_cli.py | 19 ++- lib/cli/src/crewai_cli/checkpoint_tui.py | 4 +- lib/cli/src/crewai_cli/cli.py | 23 ++- lib/cli/tests/test_checkpoint_telemetry.py | 142 ++++++++++++++++++ .../src/crewai/events/event_listener.py | 44 ++++++ .../telemetry/test_checkpoint_telemetry.py | 119 +++++++++++++++ 6 files changed, 344 insertions(+), 7 deletions(-) create mode 100644 lib/cli/tests/test_checkpoint_telemetry.py create mode 100644 lib/crewai/tests/telemetry/test_checkpoint_telemetry.py diff --git a/lib/cli/src/crewai_cli/checkpoint_cli.py b/lib/cli/src/crewai_cli/checkpoint_cli.py index ca46a32d6..65cdc0e2c 100644 --- a/lib/cli/src/crewai_cli/checkpoint_cli.py +++ b/lib/cli/src/crewai_cli/checkpoint_cli.py @@ -8,11 +8,27 @@ import json import os import re import sqlite3 -from typing import Any +from typing import Any, Literal import click +def _record_checkpoint_usage( + action: Literal[ + "list", "info", "resume", "diff", "prune", "tui", "tui_resume", "tui_fork" + ], +) -> None: + """Count a CLI action without recording checkpoint data or blocking execution.""" + try: + from crewai_core.telemetry import Telemetry + + telemetry = Telemetry() + telemetry.set_tracer() + telemetry.feature_usage_span(f"cli_usage:checkpoint_{action}") + except Exception: # noqa: S110 - telemetry must never break a command + pass + + _PLACEHOLDER_RE = re.compile(r"\{([A-Za-z_][A-Za-z0-9_\-]*)}") @@ -696,6 +712,7 @@ def prune_checkpoints( return duration: timedelta | None = _parse_duration(older_than) if older_than else None + _record_checkpoint_usage("prune") deleted: int if _is_sqlite(location): diff --git a/lib/cli/src/crewai_cli/checkpoint_tui.py b/lib/cli/src/crewai_cli/checkpoint_tui.py index 43370ba49..ae1f78f77 100644 --- a/lib/cli/src/crewai_cli/checkpoint_tui.py +++ b/lib/cli/src/crewai_cli/checkpoint_tui.py @@ -26,6 +26,7 @@ from crewai_cli.checkpoint_cli import ( _is_sqlite, _list_json, _list_sqlite, + _record_checkpoint_usage, ) @@ -699,7 +700,7 @@ class CheckpointTUI(App[_TuiResult]): if event.node.data is not None: await self._show_detail(event.node.data) - def _exit_with_action(self, action: str) -> None: + def _exit_with_action(self, action: Literal["resume", "fork"]) -> None: if self._selected_entry is None: self.notify("No checkpoint selected", severity="warning") return @@ -709,6 +710,7 @@ class CheckpointTUI(App[_TuiResult]): etype = self._detect_entity_type(self._selected_entry) name = self._selected_entry.get("name", "")[:30] self.notify(f"{action.title()}: {name}") + _record_checkpoint_usage("tui_resume" if action == "resume" else "tui_fork") self.exit((loc, action, inputs, overrides, etype)) def action_resume(self) -> None: diff --git a/lib/cli/src/crewai_cli/cli.py b/lib/cli/src/crewai_cli/cli.py index 5b5ab984e..4c922f84c 100644 --- a/lib/cli/src/crewai_cli/cli.py +++ b/lib/cli/src/crewai_cli/cli.py @@ -1268,7 +1268,7 @@ def traces_status() -> None: @click.pass_context def checkpoint(ctx: click.Context, location: str) -> None: """Browse and inspect checkpoints. Launches a TUI when called without a subcommand.""" - from crewai_cli.checkpoint_cli import _detect_location + from crewai_cli.checkpoint_cli import _detect_location, _record_checkpoint_usage location = _detect_location(location) ctx.ensure_object(dict) @@ -1276,6 +1276,7 @@ def checkpoint(ctx: click.Context, location: str) -> None: if ctx.invoked_subcommand is None: from crewai_cli.checkpoint_tui import run_checkpoint_tui + _record_checkpoint_usage("tui") run_checkpoint_tui(location) @@ -1283,8 +1284,13 @@ def checkpoint(ctx: click.Context, location: str) -> None: @click.argument("location", default="./.checkpoints") def checkpoint_list(location: str) -> None: """List checkpoints in a directory.""" - from crewai_cli.checkpoint_cli import _detect_location, list_checkpoints + from crewai_cli.checkpoint_cli import ( + _detect_location, + _record_checkpoint_usage, + list_checkpoints, + ) + _record_checkpoint_usage("list") list_checkpoints(_detect_location(location)) @@ -1292,8 +1298,13 @@ def checkpoint_list(location: str) -> None: @click.argument("path", default="./.checkpoints") def checkpoint_info(path: str) -> None: """Show details of a checkpoint. Pass a file or directory for latest.""" - from crewai_cli.checkpoint_cli import _detect_location, info_checkpoint + from crewai_cli.checkpoint_cli import ( + _detect_location, + _record_checkpoint_usage, + info_checkpoint, + ) + _record_checkpoint_usage("info") info_checkpoint(_detect_location(path)) @@ -1302,8 +1313,9 @@ def checkpoint_info(path: str) -> None: @click.pass_context def checkpoint_resume(ctx: click.Context, checkpoint_id: str | None) -> None: """Resume from a checkpoint. Defaults to the most recent.""" - from crewai_cli.checkpoint_cli import resume_checkpoint + from crewai_cli.checkpoint_cli import _record_checkpoint_usage, resume_checkpoint + _record_checkpoint_usage("resume") resume_checkpoint(ctx.obj["location"], checkpoint_id) @@ -1313,8 +1325,9 @@ def checkpoint_resume(ctx: click.Context, checkpoint_id: str | None) -> None: @click.pass_context def checkpoint_diff(ctx: click.Context, id1: str, id2: str) -> None: """Compare two checkpoints side-by-side.""" - from crewai_cli.checkpoint_cli import diff_checkpoints + from crewai_cli.checkpoint_cli import _record_checkpoint_usage, diff_checkpoints + _record_checkpoint_usage("diff") diff_checkpoints(ctx.obj["location"], id1, id2) diff --git a/lib/cli/tests/test_checkpoint_telemetry.py b/lib/cli/tests/test_checkpoint_telemetry.py new file mode 100644 index 000000000..518bd0189 --- /dev/null +++ b/lib/cli/tests/test_checkpoint_telemetry.py @@ -0,0 +1,142 @@ +"""Checkpoint CLI usage is counted once, without including checkpoint data.""" + +from unittest.mock import MagicMock, patch + +from click.testing import CliRunner +import pytest + +from crewai_cli.checkpoint_cli import _record_checkpoint_usage +from crewai_cli.checkpoint_tui import CheckpointTUI +from crewai_cli.cli import crewai +from crewai_core.telemetry import Telemetry + + +@pytest.mark.parametrize( + ("args", "target", "action"), + [ + ([], "checkpoint_tui.run_checkpoint_tui", "tui"), + (["list", "/private/checkpoints"], "checkpoint_cli.list_checkpoints", "list"), + ( + ["info", "/private/checkpoint.json"], + "checkpoint_cli.info_checkpoint", + "info", + ), + (["resume", "private-id"], "checkpoint_cli.resume_checkpoint", "resume"), + ( + ["diff", "private-id-1", "private-id-2"], + "checkpoint_cli.diff_checkpoints", + "diff", + ), + ], +) +def test_checkpoint_command_usage(args, target, action): + with ( + patch("crewai_core.telemetry.Telemetry") as telemetry, + patch(f"crewai_cli.{target}") as operation, + ): + result = CliRunner().invoke(crewai, ["checkpoint", *args]) + + assert result.exit_code == 0, result.output + operation.assert_called_once() + telemetry.return_value.feature_usage_span.assert_called_once_with( + f"cli_usage:checkpoint_{action}" + ) + + +@pytest.mark.parametrize("args", [["--help"], ["list", "--help"], ["diff"]]) +def test_help_and_invalid_arguments_are_not_counted(args): + with patch("crewai_core.telemetry.Telemetry") as telemetry: + CliRunner().invoke(crewai, ["checkpoint", *args]) + telemetry.assert_not_called() + + +@pytest.mark.parametrize( + "failure", ["initialization", "set_tracer", "feature_usage_span"] +) +def test_telemetry_failure_does_not_block_command(failure): + with ( + patch("crewai_core.telemetry.Telemetry") as telemetry, + patch("crewai_cli.checkpoint_cli.list_checkpoints") as operation, + ): + if failure == "initialization": + telemetry.side_effect = RuntimeError("telemetry unavailable") + else: + getattr(telemetry.return_value, failure).side_effect = RuntimeError( + "unavailable" + ) + result = CliRunner().invoke(crewai, ["checkpoint", "list"]) + + assert result.exit_code == 0, result.output + operation.assert_called_once() + + +@pytest.mark.parametrize( + "flag", ["OTEL_SDK_DISABLED", "CREWAI_DISABLE_TELEMETRY", "CREWAI_DISABLE_TRACKING"] +) +def test_checkpoint_usage_respects_telemetry_opt_out(monkeypatch, flag): + monkeypatch.setenv(flag, "true") + monkeypatch.setattr(Telemetry, "_instance", None) + telemetry = Telemetry() + telemetry.provider = MagicMock() + telemetry.ready = True + + _record_checkpoint_usage("list") + + telemetry.provider.get_tracer.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["resume", "fork"]) +@pytest.mark.parametrize("selected", [False, True]) +async def test_tui_counts_only_actions_with_a_selected_checkpoint(action, selected): + app = CheckpointTUI(location="/nonexistent/checkpoints") + with ( + patch("crewai_core.telemetry.Telemetry") as telemetry, + patch.object(app, "_collect_inputs", return_value={}), + patch.object(app, "_collect_task_overrides", return_value={}), + patch.object(app, "_resolve_location", return_value="/private/checkpoint.json"), + patch.object(app, "_detect_entity_type", return_value="crew"), + patch.object(app, "exit") as exit_app, + ): + async with app.run_test(): + app._selected_entry = {"name": "private-id"} if selected else None + getattr(app, f"action_{action}")() + + if selected: + telemetry.return_value.feature_usage_span.assert_called_once_with( + f"cli_usage:checkpoint_tui_{action}" + ) + exit_app.assert_called_once() + else: + telemetry.assert_not_called() + exit_app.assert_not_called() + + +@pytest.mark.parametrize("dry_run", [False, True]) +@pytest.mark.parametrize("retention", [[], ["--older-than", "invalid"]]) +def test_invalid_prune_is_not_counted(tmp_path, dry_run, retention): + args = ["checkpoint", "--location", str(tmp_path), "prune", *retention] + if dry_run: + args.append("--dry-run") + with patch("crewai_core.telemetry.Telemetry") as telemetry: + result = CliRunner().invoke(crewai, args) + if retention: + assert result.exit_code == 2 + assert "Invalid duration" in result.output + else: + assert "Specify --keep N and/or --older-than" in result.output + telemetry.assert_not_called() + + +@pytest.mark.parametrize("dry_run", [False, True]) +@pytest.mark.parametrize("retention", [["--keep", "2"], ["--older-than", "7d"]]) +def test_valid_prune_is_counted_once(tmp_path, dry_run, retention): + args = ["checkpoint", "--location", str(tmp_path), "prune", *retention] + if dry_run: + args.append("--dry-run") + with patch("crewai_core.telemetry.Telemetry") as telemetry: + result = CliRunner().invoke(crewai, args) + assert result.exit_code == 0, result.output + telemetry.return_value.feature_usage_span.assert_called_once_with( + "cli_usage:checkpoint_prune" + ) diff --git a/lib/crewai/src/crewai/events/event_listener.py b/lib/crewai/src/crewai/events/event_listener.py index fd0d77e34..94c4f6f7c 100644 --- a/lib/crewai/src/crewai/events/event_listener.py +++ b/lib/crewai/src/crewai/events/event_listener.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any from pydantic import Field, PrivateAttr from crewai.events.base_event_listener import BaseEventListener +from crewai.events.event_bus import is_replaying from crewai.events.listeners.tracing.trace_listener import TraceCollectionListener from crewai.events.types.a2a_events import ( A2AConversationCompletedEvent, @@ -23,6 +24,14 @@ from crewai.events.types.agent_events import ( LiteAgentExecutionErrorEvent, LiteAgentExecutionStartedEvent, ) +from crewai.events.types.checkpoint_events import ( + CheckpointCompletedEvent, + CheckpointFailedEvent, + CheckpointForkCompletedEvent, + CheckpointPrunedEvent, + CheckpointRestoreCompletedEvent, + CheckpointRestoreFailedEvent, +) from crewai.events.types.crew_events import ( CrewKickoffCompletedEvent, CrewKickoffFailedEvent, @@ -1000,6 +1009,41 @@ class EventListener(BaseEventListener): ) -> None: self._telemetry.feature_usage_span("memory:retrieval") + # Replayed checkpoint events describe past operations, not new usage. + @crewai_event_bus.on(CheckpointCompletedEvent) + def on_checkpoint_save(_: Any, event: CheckpointCompletedEvent) -> None: + if not is_replaying(): + self._telemetry.feature_usage_span("checkpoint:save") + + @crewai_event_bus.on(CheckpointFailedEvent) + def on_checkpoint_save_failed(_: Any, event: CheckpointFailedEvent) -> None: + if not is_replaying(): + self._telemetry.feature_usage_span("checkpoint:save_failed") + + @crewai_event_bus.on(CheckpointRestoreCompletedEvent) + def on_checkpoint_restore( + _: Any, event: CheckpointRestoreCompletedEvent + ) -> None: + if not is_replaying(): + self._telemetry.feature_usage_span("checkpoint:restore") + + @crewai_event_bus.on(CheckpointRestoreFailedEvent) + def on_checkpoint_restore_failed( + _: Any, event: CheckpointRestoreFailedEvent + ) -> None: + if not is_replaying(): + self._telemetry.feature_usage_span("checkpoint:restore_failed") + + @crewai_event_bus.on(CheckpointForkCompletedEvent) + def on_checkpoint_fork(_: Any, event: CheckpointForkCompletedEvent) -> None: + if not is_replaying(): + self._telemetry.feature_usage_span("checkpoint:fork") + + @crewai_event_bus.on(CheckpointPrunedEvent) + def on_checkpoint_prune(_: Any, event: CheckpointPrunedEvent) -> None: + if not is_replaying(): + self._telemetry.feature_usage_span("checkpoint:prune") + @crewai_event_bus.on(CrewKickoffStartedEvent) def on_crew_kickoff_hooks(_: Any, event: CrewKickoffStartedEvent) -> None: from crewai.hooks.llm_hooks import ( diff --git a/lib/crewai/tests/telemetry/test_checkpoint_telemetry.py b/lib/crewai/tests/telemetry/test_checkpoint_telemetry.py new file mode 100644 index 000000000..d504b06c6 --- /dev/null +++ b/lib/crewai/tests/telemetry/test_checkpoint_telemetry.py @@ -0,0 +1,119 @@ +"""Runtime checkpoint operations report feature usage, including Python-only calls.""" + +from collections import Counter +from unittest.mock import patch + +import pytest + +from crewai.events.event_bus import crewai_event_bus +from crewai.events.event_listener import event_listener +from crewai.events.types.checkpoint_events import ( + CheckpointCompletedEvent, + CheckpointFailedEvent, + CheckpointForkCompletedEvent, + CheckpointPrunedEvent, + CheckpointRestoreCompletedEvent, + CheckpointRestoreFailedEvent, +) +from crewai.state.checkpoint_config import CheckpointConfig +from crewai.state.checkpoint_listener import _do_checkpoint +from crewai.state.provider.json_provider import JsonProvider +from crewai.state.runtime import RuntimeState + +from ..utils import wait_for_event_handlers + + +@pytest.fixture +def features(monkeypatch): + recorded = [] + with crewai_event_bus.scoped_handlers(): + event_listener.setup_listeners(crewai_event_bus) + monkeypatch.setattr( + event_listener._telemetry, "feature_usage_span", recorded.append + ) + yield recorded + wait_for_event_handlers() + + +def test_manual_save_restore_and_fork(tmp_path, features): + state = RuntimeState(root=[]) + location = state.checkpoint(str(tmp_path)) + restored = RuntimeState.from_checkpoint(CheckpointConfig(restore_from=location)) + restored.fork("private-branch") + wait_for_event_handlers() + + assert Counter(features) == { + "checkpoint:save": 1, + "checkpoint:restore": 1, + "checkpoint:fork": 1, + } + + +@pytest.mark.asyncio +async def test_async_save_and_restore(tmp_path, features): + state = RuntimeState(root=[]) + location = await state.acheckpoint(str(tmp_path)) + await RuntimeState.afrom_checkpoint(CheckpointConfig(restore_from=location)) + wait_for_event_handlers() + + assert Counter(features) == {"checkpoint:save": 1, "checkpoint:restore": 1} + + +def test_automatic_checkpoint_and_prune(tmp_path, features): + _do_checkpoint( + RuntimeState(root=[]), + CheckpointConfig(location=str(tmp_path), max_checkpoints=1), + ) + wait_for_event_handlers() + + assert Counter(features) == {"checkpoint:save": 1, "checkpoint:prune": 1} + + +def test_failed_operations_do_not_count_as_success(tmp_path, features): + with patch.object(JsonProvider, "checkpoint", side_effect=OSError("private path")): + with pytest.raises(OSError): + RuntimeState(root=[]).checkpoint(str(tmp_path)) + with pytest.raises(FileNotFoundError): + RuntimeState.from_checkpoint( + CheckpointConfig(restore_from=str(tmp_path / "missing.json")) + ) + wait_for_event_handlers() + + assert Counter(features) == { + "checkpoint:save_failed": 1, + "checkpoint:restore_failed": 1, + } + + +@pytest.mark.parametrize( + "event", + [ + CheckpointCompletedEvent( + location="private", + provider="JsonProvider", + checkpoint_id="private-id", + duration_ms=1, + ), + CheckpointFailedEvent( + location="private", provider="JsonProvider", error="private error" + ), + CheckpointRestoreCompletedEvent( + location="private", checkpoint_id="private-id", duration_ms=1 + ), + CheckpointRestoreFailedEvent(location="private", error="private error"), + CheckpointForkCompletedEvent(branch="private-branch"), + CheckpointPrunedEvent( + location="private", + provider="JsonProvider", + removed_count=1, + max_checkpoints=1, + ), + ], +) +def test_replayed_checkpoint_events_are_not_counted(event, features): + future = crewai_event_bus.replay(None, event) + if future is not None: + future.result(timeout=5) + wait_for_event_handlers() + + assert features == []