ensure_agent_sdk.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234
  1. #!/usr/bin/env python3
  2. """SessionStart bootstrap: ensure claude_agent_sdk is importable for the
  3. agentic commit reviewer.
  4. If claude_agent_sdk already imports in the current python3, this is a no-op.
  5. Otherwise it creates a venv at ~/.claude/security/agent-sdk-venv and installs
  6. the SDK there. security_reminder_hook.py prepends that venv's site-packages to
  7. sys.path before attempting the SDK import, so the venv is used as a
  8. fallback only when the system install is missing.
  9. The venv lives under ~/.claude/security/ (same dir the plugin already uses
  10. for per-session state) so it persists across plugin updates — rebuilding
  11. on every update is 30-60s of wasted work for a package that changes far
  12. less often than the plugin does.
  13. """
  14. from __future__ import annotations
  15. import importlib.util
  16. import json
  17. import os
  18. import subprocess
  19. import sys
  20. import time
  21. from pathlib import Path
  22. # Outcome codes for the sdk_bootstrap metric. Values are stable for telemetry.
  23. NOOP_SYSTEM = 0 # claude_agent_sdk already importable in system python
  24. NOOP_VENV = 1 # venv already built and SDK imports from it
  25. BUILT = 2 # venv created + SDK pip-installed this run
  26. BUILD_FAILED = 3 # venv create or pip install raised/timed out
  27. # Outcome 4 was previously SKIP_WIN32; retired now that the consumer glob in
  28. # llm.py also matches Windows venv layout (Lib/site-packages). Don't reuse the
  29. # value — telemetry rows from older plugin builds still emit 4.
  30. SKIP_SENTINEL = 5 # another SessionStart is currently building
  31. def _sdk_on_syspath() -> bool:
  32. # find_spec is ~10ms; actually importing the SDK pulls in
  33. # transitive deps and costs ~800ms — too heavy for a
  34. # per-SessionStart no-op check that most sessions hit.
  35. try:
  36. return importlib.util.find_spec("claude_agent_sdk") is not None
  37. except Exception:
  38. return False
  39. def _plugin_version_int() -> int:
  40. # Same encoding as security_reminder_hook._read_plugin_version_int so
  41. # metrics rows from both hooks join on pv.
  42. try:
  43. p = Path(__file__).parent.parent / ".claude-plugin" / "plugin.json"
  44. v = json.loads(p.read_text())["version"]
  45. major, minor, patch = (int(x) for x in v.split(".")[:3])
  46. return major * 10000 + minor * 100 + patch
  47. except Exception:
  48. return 0
  49. def main() -> tuple[int, str, str]:
  50. """Run the bootstrap. Returns (outcome, err_phase, err_kind).
  51. err_phase / err_kind are non-empty only on BUILD_FAILED — they let
  52. telemetry split bootstrap failures by root cause.
  53. """
  54. if _sdk_on_syspath():
  55. return NOOP_SYSTEM, "", ""
  56. state_dir = Path(
  57. os.environ.get("SECURITY_WARNINGS_STATE_DIR")
  58. or os.path.expanduser("~/.claude/security")
  59. )
  60. venv = state_dir / "agent-sdk-venv"
  61. # Windows venvs put the interpreter at Scripts\python.exe; POSIX uses bin/python.
  62. if sys.platform == "win32":
  63. venv_py = venv / "Scripts" / "python.exe"
  64. else:
  65. venv_py = venv / "bin" / "python"
  66. # Another SessionStart (concurrent CC instance, same plugin) may already
  67. # be building. The sentinel lives NEXT TO the venv, not inside it —
  68. # `python -m venv --clear` wipes the target dir's contents, so an
  69. # in-venv sentinel would be deleted the instant we create the venv.
  70. # Stale sentinels (>5min) from a SIGKILL'd build are ignored.
  71. sentinel = state_dir / "agent-sdk-venv.building"
  72. if sentinel.exists():
  73. try:
  74. if time.time() - sentinel.stat().st_mtime < 300:
  75. return SKIP_SENTINEL, "", ""
  76. sentinel.unlink(missing_ok=True)
  77. except OSError:
  78. return SKIP_SENTINEL, "", ""
  79. # If a venv already exists and its python can import the SDK, done.
  80. if venv_py.exists():
  81. try:
  82. r = subprocess.run(
  83. [str(venv_py), "-c", "import claude_agent_sdk"],
  84. capture_output=True, timeout=10,
  85. )
  86. if r.returncode == 0:
  87. return NOOP_VENV, "", ""
  88. except Exception:
  89. pass # broken venv; rebuild below
  90. err_phase = ""
  91. err_kind = ""
  92. we_own_sentinel = False
  93. try:
  94. state_dir.mkdir(parents=True, exist_ok=True)
  95. # O_EXCL makes the sentinel an atomic lock — if two SessionStarts
  96. # race past the exists() check above, only one creates it.
  97. try:
  98. os.close(os.open(sentinel, os.O_CREAT | os.O_EXCL | os.O_WRONLY))
  99. except FileExistsError:
  100. return SKIP_SENTINEL, "", ""
  101. we_own_sentinel = True
  102. err_phase = "venv"
  103. subprocess.run(
  104. [sys.executable, "-m", "venv", "--clear", str(venv)],
  105. capture_output=True, timeout=60, check=True,
  106. )
  107. # Some machines route pip through a private registry; we
  108. # don't pass --index-url here so we inherit that default. Outside
  109. # the user's machine, pip's own default registry applies — that's the same
  110. # exposure the user would have running `pip install` themselves, so
  111. # we're not widening the supply-chain surface.
  112. #
  113. # --prefer-binary: on ARM64 Windows, pip's default resolver picks a
  114. # `cryptography` version with no published binary wheel and tries to
  115. # build from source, which needs Rust/Cargo (almost never present
  116. # on user machines). The build fails and the whole bootstrap returns
  117. # BUILD_FAILED. A binary wheel exists on PyPI for an adjacent
  118. # version (`cryptography-46.0.3-cp311-abi3-win_arm64.whl`);
  119. # --prefer-binary tells pip to pick it. Cross-platform safe: no-op
  120. # on platforms where the latest version already has a wheel.
  121. err_phase = "pip"
  122. subprocess.run(
  123. [str(venv_py), "-m", "pip", "install", "--quiet",
  124. "--disable-pip-version-check", "--prefer-binary",
  125. "claude-agent-sdk"],
  126. capture_output=True, timeout=120, check=True,
  127. )
  128. return BUILT, "", ""
  129. except subprocess.CalledProcessError as e:
  130. # Capture a stderr fingerprint so telemetry can split BUILD_FAILED by
  131. # root cause (no-network, package-not-found, dns-fail, etc.).
  132. # Categorize first, then keep a short raw tail for the long tail of
  133. # unexpected modes.
  134. stderr_b = e.stderr or b""
  135. if isinstance(stderr_b, bytes):
  136. stderr_str = stderr_b.decode("utf-8", errors="replace")
  137. else:
  138. stderr_str = str(stderr_b)
  139. s = stderr_str.lower()
  140. if "no matching distribution" in s or "could not find a version" in s:
  141. err_kind = "pip_no_match"
  142. elif "name or service not known" in s or "name resolution" in s \
  143. or "nodename nor servname" in s or "temporary failure in name" in s:
  144. err_kind = "dns_fail"
  145. elif "connection refused" in s or "connection reset" in s:
  146. err_kind = "conn_refused"
  147. elif "ssl" in s and ("verify" in s or "certificate" in s):
  148. err_kind = "ssl_verify"
  149. elif "permission denied" in s or "read-only file system" in s:
  150. err_kind = "perm_denied"
  151. elif "no module named pip" in s or "no module named ensurepip" in s:
  152. err_kind = "no_pip"
  153. elif "no space left" in s or "disk quota" in s:
  154. err_kind = "disk_full"
  155. elif "proxy" in s and ("authent" in s or "tunnel" in s or "407" in s):
  156. err_kind = "proxy_auth"
  157. elif "timeout" in s or "timed out" in s:
  158. err_kind = "stderr_timeout"
  159. else:
  160. # First 60 chars of the last non-empty stderr line — bounded to
  161. # stay inside CC's metric value-length budget. Real failure modes
  162. # we haven't categorized show up here as a low-cardinality bucket.
  163. tail = next(
  164. (ln.strip() for ln in reversed(stderr_str.splitlines()) if ln.strip()),
  165. "",
  166. )[:60]
  167. err_kind = f"other:{tail}" if tail else "other"
  168. return BUILD_FAILED, err_phase, err_kind
  169. except subprocess.TimeoutExpired:
  170. return BUILD_FAILED, err_phase, "subprocess_timeout"
  171. except Exception as e:
  172. return BUILD_FAILED, err_phase, f"exc:{type(e).__name__}"
  173. finally:
  174. # Only remove the sentinel if THIS process created it. The
  175. # FileExistsError path above means another process owns the lock;
  176. # unconditionally unlinking here would delete its sentinel and let
  177. # a third concurrent SessionStart `venv --clear` over the in-flight
  178. # build.
  179. if we_own_sentinel:
  180. sentinel.unlink(missing_ok=True)
  181. if __name__ == "__main__":
  182. # Tell the harness this is async — venv create + pip install can take
  183. # 30-60s on a cold cache, well past the default sync hook timeout.
  184. # SessionStart runs before the user's first prompt; doing this in the
  185. # background means the first commit-review of the session usually finds
  186. # the venv ready.
  187. print(json.dumps({"async": True, "asyncTimeout": 180000}), flush=True)
  188. t0 = time.perf_counter()
  189. try:
  190. outcome, err_phase, err_kind = main()
  191. except Exception as exc:
  192. outcome, err_phase, err_kind = (
  193. BUILD_FAILED, "main", f"exc:{type(exc).__name__}"
  194. )
  195. # CC's async-hook registry scans stdout line-by-line after process exit
  196. # and takes the FIRST non-{"async":...} JSON line as the hook response;
  197. # its `metrics` key is forwarded to the hook metrics event on the
  198. # next attachments pass. Must be a single line — the registry splits on
  199. # \n and json-parses each independently. Values must be bool|number OR
  200. # short strings (CC accepts string metric values if they're not
  201. # null). Stay inside the 10-key emit cap.
  202. metrics: dict[str, object] = {
  203. "sdk_bootstrap": outcome,
  204. "sdk_bootstrap_ms": round((time.perf_counter() - t0) * 1000),
  205. }
  206. if err_kind:
  207. # Truncate defensively; categorized values are <40 chars but the
  208. # `other:<tail>` mode could be longer. err_phase may be empty for
  209. # pre-venv failures (state_dir.mkdir perm-denied, sentinel O_EXCL
  210. # raising a non-FileExistsError OSError) — emit as "pre" so the
  211. # err_kind isn't silently dropped.
  212. metrics["sdk_bootstrap_phase"] = (err_phase or "pre")[:16]
  213. metrics["sdk_bootstrap_err"] = err_kind[:96]
  214. pv = _plugin_version_int()
  215. if pv:
  216. metrics["pv"] = pv
  217. print(json.dumps({"metrics": metrics}), flush=True)