From a17f99b2ddf3a3c6e39ff00affed03e2990ffd24 Mon Sep 17 00:00:00 2001
From: Daniel D <d.drizhuk@gmail.com>
Date: Tue, 4 Aug 2026 13:45:58 +0200
Subject: [PATCH 1/2] feat(cli): set AI_AGENT for subprocesses

---
 src/kimi_cli/__main__.py           |  4 ++++
 tests/core/test_startup_imports.py | 24 ++++++++++++++++++++++++
 2 files changed, 28 insertions(+)

diff --git a/src/kimi_cli/__main__.py b/src/kimi_cli/__main__.py
index 2d8aded024..af83905668 100644
--- a/src/kimi_cli/__main__.py
+++ b/src/kimi_cli/__main__.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import os
 import sys
 from collections.abc import Sequence
 from pathlib import Path
@@ -10,6 +11,9 @@ def _prog_name() -> str:
 
 
 def main(argv: Sequence[str] | None = None) -> int | str | None:
+    if not os.environ.get("AI_AGENT", "").strip():
+        os.environ["AI_AGENT"] = "kimi"
+
     from kimi_cli.telemetry.crash import install_crash_handlers, set_phase
     from kimi_cli.utils.proxy import normalize_proxy_env
 
diff --git a/tests/core/test_startup_imports.py b/tests/core/test_startup_imports.py
index 538d66524b..ef1b5b7c53 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(
         """

From bf2084b01e0fb971ec97374d0abd017109b0efa8 Mon Sep 17 00:00:00 2001
From: Daniel D <d.drizhuk@gmail.com>
Date: Tue, 4 Aug 2026 14:36:10 +0200
Subject: [PATCH 2/2] fix(cli): mark standalone binary subprocesses

---
 src/kimi_cli/__main__.py             |  7 +++----
 src/kimi_cli/cli/__main__.py         |  3 +++
 src/kimi_cli/utils/subprocess_env.py |  6 ++++++
 tests/core/test_startup_imports.py   | 30 ++++++++++++++++++++++++++++
 4 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/src/kimi_cli/__main__.py b/src/kimi_cli/__main__.py
index af83905668..3319f63a05 100644
--- a/src/kimi_cli/__main__.py
+++ b/src/kimi_cli/__main__.py
@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import os
 import sys
 from collections.abc import Sequence
 from pathlib import Path
@@ -11,11 +10,11 @@ def _prog_name() -> str:
 
 
 def main(argv: Sequence[str] | None = None) -> int | str | None:
-    if not os.environ.get("AI_AGENT", "").strip():
-        os.environ["AI_AGENT"] = "kimi"
-
     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 ef1b5b7c53..e441560974 100644
--- a/tests/core/test_startup_imports.py
+++ b/tests/core/test_startup_imports.py
@@ -228,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"