diff --git a/src/kimi_cli/__main__.py b/src/kimi_cli/__main__.py index 2d8aded024..3319f63a05 100644 --- a/src/kimi_cli/__main__.py +++ b/src/kimi_cli/__main__.py @@ -12,6 +12,9 @@ def _prog_name() -> str: def main(argv: Sequence[str] | None = None) -> int | str | None: from kimi_cli.telemetry.crash import install_crash_handlers, set_phase from kimi_cli.utils.proxy import normalize_proxy_env + from kimi_cli.utils.subprocess_env import ensure_ai_agent_env + + ensure_ai_agent_env() # Install excepthook before anything else so startup-phase crashes are captured. install_crash_handlers() diff --git a/src/kimi_cli/cli/__main__.py b/src/kimi_cli/cli/__main__.py index 83fead793c..6bca0bba29 100644 --- a/src/kimi_cli/cli/__main__.py +++ b/src/kimi_cli/cli/__main__.py @@ -10,6 +10,9 @@ def main(argv: Sequence[str] | None = None) -> int | str | None: from kimi_cli.telemetry.crash import install_crash_handlers, set_phase from kimi_cli.utils.environment import GitBashNotFoundError from kimi_cli.utils.proxy import normalize_proxy_env + from kimi_cli.utils.subprocess_env import ensure_ai_agent_env + + ensure_ai_agent_env() # Same entry treatment as kimi_cli.__main__: install excepthook before # anything else so startup-phase crashes in subcommand subprocesses diff --git a/src/kimi_cli/utils/subprocess_env.py b/src/kimi_cli/utils/subprocess_env.py index 1aaf4eda49..58e6b29dc0 100644 --- a/src/kimi_cli/utils/subprocess_env.py +++ b/src/kimi_cli/utils/subprocess_env.py @@ -20,6 +20,12 @@ ] +def ensure_ai_agent_env() -> None: + """Identify Kimi CLI to subprocesses unless a wrapper already did so.""" + if not os.environ.get("AI_AGENT", "").strip(): + os.environ["AI_AGENT"] = "kimi" + + def get_clean_env(base_env: dict[str, str] | None = None) -> dict[str, str]: """ Get a clean environment suitable for spawning subprocesses. diff --git a/tests/core/test_startup_imports.py b/tests/core/test_startup_imports.py index 538d66524b..e441560974 100644 --- a/tests/core/test_startup_imports.py +++ b/tests/core/test_startup_imports.py @@ -153,6 +153,30 @@ def test_package_entrypoint_fast_path_avoids_cli_import() -> None: assert proc.stdout.strip() == "ok" +def test_package_entrypoint_sets_ai_agent_marker() -> None: + proc = _run_python( + """ +import io +import os +from contextlib import redirect_stdout + +from kimi_cli.__main__ import main + +os.environ.pop("AI_AGENT", None) +with redirect_stdout(io.StringIO()): + main(["--version"]) +assert os.environ["AI_AGENT"] == "kimi" + +os.environ["AI_AGENT"] = "wrapper" +with redirect_stdout(io.StringIO()): + main(["--version"]) +assert os.environ["AI_AGENT"] == "wrapper" +print("ok") +""" + ) + assert proc.stdout.strip() == "ok" + + def test_package_entrypoint_falls_back_to_cli_for_commands() -> None: proc = _run_python( """ @@ -204,3 +228,33 @@ def fake_cli(*_args, **_kwargs): """ ) assert proc.stdout.strip() == "ok" + + +def test_cli_module_entrypoint_sets_ai_agent_marker() -> None: + proc = _run_python( + """ +import os + +import kimi_cli.cli.__main__ as cli_main + +seen = [] + +def fake_cli(*_args, **_kwargs): + seen.append(os.environ["AI_AGENT"]) + +cli_main.cli = fake_cli + +os.environ.pop("AI_AGENT", None) +cli_main.main([]) + +os.environ["AI_AGENT"] = " " +cli_main.main([]) + +os.environ["AI_AGENT"] = "wrapper" +cli_main.main([]) + +assert seen == ["kimi", "kimi", "wrapper"] +print("ok") +""" + ) + assert proc.stdout.strip() == "ok"