feat(cli): add global CPU profiling (#5141)

## Related issue

[OMNI-4272](https://linear.app/omnigent/issue/OMNI-4272/add-a-global-profile-option-to-omni-cli-for-profiling)

## Summary

- Add a global `omni --profile COMMAND` flag so developers can profile any CLI command with Python's built-in `cProfile`.
- Print actionable Omnigent-only cumulative and self-time summaries while preserving the complete timestamped `.prof` data under `~/.omnigent/profiles/`.
- Keep the existing `omni run --profile NAME` Databricks credential option compatible, including when both profile options are used together.

**ELI5:** Put `--profile` before a command to see which Omnigent functions made it slow; open the saved `.prof` file when deeper analysis is needed.

```text
omni --profile COMMAND
        |
        v
     cProfile
      /    \
focused stderr  full timestamped .prof
```

## Test Plan

- `uv run ruff format --check omnigent/cli.py tests/cli/test_cli.py`
- `uv run ruff check omnigent/cli.py tests/cli/test_cli.py`
- `uv run --group test pytest -q tests/cli/test_cli.py::test_global_profile_writes_summary_and_timestamped_stats tests/cli/test_cli.py::test_removed_ad_hoc_detection tests/cli/test_cli.py::test_help_groups_harnesses_and_other_commands tests/cli/test_cli.py::test_run_profile_sets_databricks_config_profile_env tests/cli/test_cli.py::test_global_cpu_profile_coexists_with_run_databricks_profile tests/cli/test_cli.py::test_run_profile_wins_over_preset_env tests/cli/test_cli.py::test_run_without_profile_leaves_preset_env_untouched`
- Manually ran `OMNIGENT_DATA_DIR=<temp-dir> uv run python -m omnigent --profile version` and inspected the focused summary and generated `.prof` file.

## Demo

N/A

## Type of change

- [ ] Bug fix
- [x] Feature
- [ ] UI / frontend change
- [ ] Refactor / chore
- [ ] Docs
- [ ] Test / CI
- [ ] Breaking change

## Test coverage

- [x] Unit tests added / updated
- [ ] Integration tests added / updated
- [ ] E2E tests added / updated
- [x] Manual verification completed
- [ ] Existing tests cover this change
- [ ] Not applicable

## Coverage notes

Verified the real `python -m omnigent` entry point creates a loadable timestamped profile, renders only actionable Omnigent functions in the summary, and preserves the separate Databricks `run --profile NAME` behavior.

## Changelog

Use `omni --profile COMMAND` to print focused CPU profiling results and save the full profile for deeper analysis.

Signed-off-by: Zeyi (Rice) Fan <zeyi.f@databricks.com>

omnigent/cli.py

@@ -1781,6 +1781,98 @@ def _set_log_to_stderr(
     return value
 
 
+def _finish_cli_profile(profiler: Any, output_path: Path) -> None:  # type: ignore[explicit-any]
+    """Stop a requested cProfile run and render a focused bottleneck summary."""
+    import pstats
+
+    profiler.disable()
+    profiler.dump_stats(output_path)
+    stats = pstats.Stats(profiler)
+    stats_state = vars(stats)
+    raw_stats = cast(
+        dict[tuple[str, int, str], tuple[int, int, float, float, object]],
+        stats_state["stats"],
+    )
+    total_time = float(stats_state["total_tt"])
+    total_calls = int(stats_state["total_calls"])
+    primitive_calls = int(stats_state["prim_calls"])
+    package_root = Path(__file__).resolve().parent
+    source_root = package_root.parent
+
+    # (filename, line, function, primitive calls, total calls, self, cumulative)
+    rows = [(*key, cc, nc, tt, ct) for key, (cc, nc, tt, ct, _) in raw_stats.items()]
+
+    def _location(filename: str, line: int, function: str) -> str:
+        path = Path(filename).resolve()
+        try:
+            display = str(path.relative_to(source_root))
+        except ValueError:
+            display = path.name
+        return f"{display}:{line}({function})"
+
+    def _duration(seconds: float) -> str:
+        if seconds < 0.01:
+            return f"{seconds * 1_000:.2f} ms"
+        if seconds < 1:
+            return f"{seconds * 1_000:.1f} ms"
+        return f"{seconds:.3f} s"
+
+    def _print_rows(
+        title: str,
+        selected: list[tuple[str, int, str, int, int, float, float]],
+    ) -> None:
+        click.echo(f"\n{title}", err=True)
+        click.echo(f"  {'self':>9}  {'cumulative':>10}  {'calls':>9}  function", err=True)
+        for filename, line, function, primitive, calls, self_time, cumulative in selected[:10]:
+            call_count = str(calls) if calls == primitive else f"{calls}/{primitive}"
+            click.echo(
+                f"  {_duration(self_time):>9}  {_duration(cumulative):>10}  "
+                f"{call_count:>9}  {_location(filename, line, function)}",
+                err=True,
+            )
+
+    omnigent_rows = [
+        row for row in rows if Path(row[0]).resolve().is_relative_to(package_root) and row[6] > 0
+    ]
+    omnigent_rows.sort(key=lambda row: row[6], reverse=True)
+    self_time_rows = [row for row in omnigent_rows if row[5] > 0]
+    self_time_rows.sort(key=lambda row: row[5], reverse=True)
+
+    click.echo(
+        f"\nCLI profile: {_duration(total_time)}, "
+        f"{total_calls:,} calls ({primitive_calls:,} primitive)",
+        err=True,
+    )
+    _print_rows("Top Omnigent call paths", omnigent_rows)
+    _print_rows("Top Omnigent functions by self time", self_time_rows)
+    click.echo(f"\nFull profile data: {output_path}", err=True)
+
+
+def _start_cli_profile(
+    ctx: click.Context,
+    _param: click.Parameter,
+    value: bool,
+) -> bool:
+    """Start cProfile early and finish it after the selected command exits."""
+    if not value:
+        return value
+
+    import cProfile
+
+    timestamp = time.strftime("%Y%m%d-%H%M%S")
+    microseconds = time.time_ns() // 1_000 % 1_000_000
+    profile_dir = data_dir() / "profiles"
+    profile_dir.mkdir(parents=True, exist_ok=True)
+    output_path = profile_dir / (f"omnigent-cli-{timestamp}-{microseconds:06d}-{os.getpid()}.prof")
+    profiler = cProfile.Profile()
+    profiler.enable()
+    ctx.call_on_close(lambda: _finish_cli_profile(profiler, output_path))
+    # Click closes callbacks last-in-first-out, so measurement stops before
+    # rendering the summary above.
+    ctx.call_on_close(profiler.disable)
+    return value
+
+
 def _extract_global_logging_flags(argv: list[str]) -> tuple[list[str], bool, bool]:
     """Remove global logging flags before run-shorthand rewriting."""
     debug_logging = False
@@ -1801,6 +1893,17 @@ def _extract_global_logging_flags(argv: list[str]) -> tuple[list[str], bool, boo
 
 
 @click.group(cls=_OmnigentCLI)
+@click.option(
+    "--profile",
+    is_flag=True,
+    is_eager=True,
+    expose_value=False,
+    callback=_start_cli_profile,
+    help=(
+        "Profile CLI execution, print a summary, and write a timestamped .prof file. "
+        "Place before COMMAND."
+    ),
+)
 @click.option(
     "--debug",
     "debug_logging",
@@ -2031,7 +2134,17 @@ def main() -> None:
     # intentionally tiny (currently only help/version); runner flags live on
     # ``run``. Treat a leading non-top-level flag as bare-run shorthand so
     # users can type the natural no-AGENT launcher form.
-    if argv and argv[0].startswith("-") and argv[0] not in {"--help", "-h", "--version"}:
+    if (
+        argv
+        and argv[0].startswith("-")
+        and argv[0]
+        not in {
+            "--help",
+            "-h",
+            "--version",
+            "--profile",
+        }
+    ):
         argv = ["run", *argv]
 
     # Shorthand: ``omnigent myagent.yaml [opts]`` → ``run myagent.yaml [opts]``.
@@ -2201,6 +2314,10 @@ def _is_removed_ad_hoc_invocation(argv: list[str]) -> bool:
     # help listing subcommands, not the legacy argparse help.
     if argv[0] in {"--help", "-h", "--version"}:
         return False
+    # A root profiling flag may precede an eager help/version flag or stand
+    # alone. These are valid Click invocations, not removed ad-hoc chat.
+    if all(token in {"--profile", "--help", "-h", "--version"} for token in argv):
+        return False
     # Skip leading flags to find the first positional. If all
     # tokens are flags (e.g. ``omnigent --system-prompt "..."``),
     # treat it as removed ad-hoc chat rather than handing it to click

tests/cli/test_cli.py

@@ -120,6 +120,39 @@ def _restore_logging_state() -> Iterator[None]:
         logger.propagate = propagate
 
 
+def test_global_profile_writes_summary_and_timestamped_stats(tmp_path: Path) -> None:
+    """The real entry point profiles any command selected after the root flag."""
+    repo_root = Path(__file__).resolve().parents[2]
+    pythonpath = os.pathsep.join(
+        part for part in (str(repo_root), os.environ.get("PYTHONPATH")) if part
+    )
+    result = subprocess.run(
+        [sys.executable, "-m", "omnigent", "--profile", "version"],
+        capture_output=True,
+        text=True,
+        timeout=30,
+        cwd=tmp_path,
+        env={
+            **os.environ,
+            "PYTHONPATH": pythonpath,
+            "OMNIGENT_DATA_DIR": str(tmp_path / "data"),
+        },
+    )
+
+    assert result.returncode == 0, result.stderr
+    assert "CLI profile:" in result.stderr
+    assert "Top Omnigent call paths" in result.stderr
+    assert "Top Omnigent functions by self time" in result.stderr
+    self_time_table = result.stderr.split("Top Omnigent functions by self time", 1)[1]
+    assert "<built-in method" not in self_time_table
+    assert "Full profile data:" in result.stderr
+    [profile_path] = list((tmp_path / "data" / "profiles").glob("omnigent-cli-*.prof"))
+
+    import pstats
+
+    assert pstats.Stats(str(profile_path)).total_calls > 0
+
+
 def test_python_module_entrypoint_uses_unified_click_cli() -> None:
     """
     ``python -m omnigent`` must dispatch through the same click CLI
@@ -191,6 +224,7 @@ def test_wrapper_guard_bypass_reaches_cli_end_to_end() -> None:
         (["run", "tests/resources/examples/hello_world.yaml"], False),
         (["attach", "tests/resources/examples/hello_world.yaml"], False),
         (["--help"], False),
+        (["--profile", "--help"], False),
         (["what does this repo do?"], True),
         (["--system-prompt", "You are terse"], True),
         # A single command-shaped word is an unknown subcommand, not
@@ -1042,10 +1076,11 @@ def test_kiro_command_is_registered_in_click_help() -> None:
 
 
 def test_help_groups_harnesses_and_other_commands() -> None:
-    """``--help`` lists a ``Harnesses`` section separate from ``Commands``."""
+    """``--help`` lists global options and separates command categories."""
     result = CliRunner().invoke(cli, ["--help"])
 
     assert result.exit_code == 0, result.output
+    assert "--profile" in result.output
     assert "Harnesses:" in result.output
     assert "Commands:" in result.output
     # A harness launcher lands under Harnesses; a management command
@@ -3458,6 +3493,36 @@ def test_run_profile_sets_databricks_config_profile_env(
     assert seen["value"] == "my-sp"
 
 
+def test_global_cpu_profile_coexists_with_run_databricks_profile(
+    monkeypatch: pytest.MonkeyPatch,
+    tmp_path: Path,
+) -> None:
+    """Root profiling and ``run --profile NAME`` keep distinct semantics."""
+    monkeypatch.setenv("OMNIGENT_DATA_DIR", str(tmp_path / "data"))
+    monkeypatch.delenv("DATABRICKS_CONFIG_PROFILE", raising=False)
+    seen = _capture_profile_env_at_dispatch(monkeypatch)
+
+    result = CliRunner().invoke(
+        cli,
+        [
+            "--profile",
+            "run",
+            "--server",
+            "https://example.com",
+            "--profile",
+            "my-sp",
+            "-p",
+            "hi",
+        ],
+    )
+
+    assert result.exit_code == 0, result.output
+    assert seen["value"] == "my-sp"
+    assert "Top Omnigent call paths" in result.stderr
+    profiles = tmp_path / "data" / "profiles"
+    assert len(list(profiles.glob("omnigent-cli-*.prof"))) == 1
+
+
 def test_run_profile_wins_over_preset_env(
     monkeypatch: pytest.MonkeyPatch,
 ) -> None: